Added disconnect and connect when already connected service code

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -93,11 +93,21 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <param name="sqlConnection"></param>
         public delegate Task OnConnectionHandler(ConnectionInfo info);
 
+        /// <summary>
+        // Callback for ondisconnect handler
+        /// </summary>
+        public delegate Task OnDisconnectHandler(ConnectionSummary summary);
+
         /// <summary>
         /// List of onconnection handlers
         /// </summary>
         private readonly List<OnConnectionHandler> onConnectionActivities = new List<OnConnectionHandler>();
 
+        /// <summary>
+        /// List of ondisconnect handlers
+        /// </summary>
+        private readonly List<OnDisconnectHandler> onDisconnectActivities = new List<OnDisconnectHandler>();
+
         /// <summary>
         /// Gets the SQL connection factory instance
         /// </summary>
@@ -128,16 +138,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             return this.ownerToConnectionMap.TryGetValue(ownerUri, out connectionInfo);
         }
 
-        private static ConnectionSummary CopySummary(ConnectionSummary summary)
-        {
-            return new ConnectionSummary()
-            {
-                ServerName = summary.ServerName,
-                DatabaseName = summary.DatabaseName,
-                UserName = summary.UserName
-            };
-        }
-
         /// <summary>
         /// Open a connection with the specified connection details
         /// </summary>
@@ -153,10 +153,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 };
             }
 
+            // Resolve if it is an existing connection
+            // Disconnect active connection if the URI is already connected
             ConnectionInfo connectionInfo;
             if (ownerToConnectionMap.TryGetValue(connectionParams.OwnerUri, out connectionInfo) )
             {
-                // TODO disconnect
+                var disconnectParams = new DisconnectParams()
+                {
+                    OwnerUri = connectionParams.OwnerUri
+                };
+                Disconnect(disconnectParams);
             }
             connectionInfo = new ConnectionInfo(ConnectionFactory, connectionParams.OwnerUri, connectionParams.Connection);
 
@@ -185,10 +191,45 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             return response;
         }
 
+        /// <summary>
+        /// Close a connection with the specified connection details.
+        /// </summary>
+        public bool Disconnect(DisconnectParams disconnectParams)
+        {
+            // Validate parameters
+            if (disconnectParams == null || String.IsNullOrEmpty(disconnectParams.OwnerUri))
+            {
+                return false;
+            }
+
+            // Lookup the connection owned by the URI
+            ConnectionInfo info;
+            if (!ownerToConnectionMap.TryGetValue(disconnectParams.OwnerUri, out info))
+            {
+                return false;
+            }
+
+            // Close the connection            
+            info.SqlConnection.Close();
+
+            // Remove URI mapping
+            ownerToConnectionMap.Remove(disconnectParams.OwnerUri);
+
+            // Invoke callback notifications
+            foreach (var activity in this.onDisconnectActivities)
+            {
+                activity(info.ConnectionDetails);
+            }
+
+            // Success
+            return true;
+        }
+
         public void InitializeService(IProtocolEndpoint serviceHost)
         {
             // Register request and event handlers with the Service Host
             serviceHost.SetRequestHandler(ConnectionRequest.Type, HandleConnectRequest);
+            serviceHost.SetRequestHandler(DisconnectRequest.Type, HandleDisconnectRequest);
 
             // Register the configuration update handler
             WorkspaceService<SqlToolsSettings>.Instance.RegisterConfigChangeCallback(HandleDidChangeConfigurationNotification);
@@ -203,6 +244,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             onConnectionActivities.Add(activity); 
         }
 
+        /// <summary>
+        /// Add a new method to be called when the ondisconnect request is submitted
+        /// </summary>
+        public void RegisterOnDisconnectTask(OnDisconnectHandler activity)
+        {
+            onDisconnectActivities.Add(activity);
+        }
+        
         /// <summary>
         /// Handle new connection requests
         /// </summary>
@@ -227,6 +276,27 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             }
         }
 
+        /// <summary>
+        /// Handle disconnect requests
+        /// </summary>
+        protected async Task HandleDisconnectRequest(
+            DisconnectParams disconnectParams,
+            RequestContext<bool> requestContext)
+        {
+            Logger.Write(LogLevel.Verbose, "HandleDisconnectRequest");
+
+            try
+            {
+                bool result = ConnectionService.Instance.Disconnect(disconnectParams);
+                await requestContext.SendResult(result);
+            }
+            catch(Exception ex)
+            {
+                await requestContext.SendError(ex.Message);
+            }
+
+        }
+        
         public Task HandleDidChangeConfigurationNotification(
             SqlToolsSettings newSettings, 
             SqlToolsSettings oldSettings, 

src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/ConnectionMessages.cs

@@ -35,7 +35,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
         /// A URI identifying the owner of the connection. This will most commonly be a file in the workspace
         /// or a virtual file representing an object in a database.         
         /// </summary>
-        public string ownerUri { get; set; }
+        public string OwnerUri { get; set; }
     }
 
     /// <summary>
@@ -47,7 +47,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
         /// A URI identifying the owner of the connection. This will most commonly be a file in the workspace
         /// or a virtual file representing an object in a database.         
         /// </summary>
-        public string ownerUri { get; set; }
+        public string OwnerUri { get; set; }
         /// <summary>
         /// Contains the high-level properties about the connection, for display to the user.
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/AutoCompleteService.cs

@@ -21,8 +21,13 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         // connection used to query for intellisense info
         private DbConnection connection;
 
+        // number of documents (URI's) that are using the cache for the same database
+        // the autocomplete service uses this to remove unreferenced caches
+        public int ReferenceCount { get; set; }
+
         public IntellisenseCache(ISqlConnectionFactory connectionFactory, ConnectionDetails connectionDetails)
         {
+            ReferenceCount = 0;
             DatabaseInfo = CopySummary(connectionDetails);
 
             // TODO error handling on this. Intellisense should catch or else the service should handle
@@ -201,6 +206,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         // Dictionary of unique intellisense caches for each Connection
         private Dictionary<ConnectionSummary, IntellisenseCache> caches = 
             new Dictionary<ConnectionSummary, IntellisenseCache>(new ConnectionSummaryComparer());
+        private Object cachesLock = new Object(); // Used when we insert/remove something from the cache dictionary
 
         private ISqlConnectionFactory factory;
         private Object factoryLock = new Object();
@@ -233,6 +239,9 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         {
             // Register a callback for when a connection is created
             ConnectionService.Instance.RegisterOnConnectionTask(UpdateAutoCompleteCache);
+
+            // Register a callback for when a connection is closed
+            ConnectionService.Instance.RegisterOnDisconnectTask(RemoveAutoCompleteCacheUriReference);
         }
 
         private async Task UpdateAutoCompleteCache(ConnectionInfo connectionInfo)
@@ -243,18 +252,48 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             }
         }
 
+        /// <summary>
+        /// Remove a reference to an autocomplete cache from a URI. If
+        /// it is the last URI connected to a particular connection,
+        /// then remove the cache.
+        /// </summary>
+        public async Task RemoveAutoCompleteCacheUriReference(ConnectionSummary summary)
+        {
+            await Task.Run( () => 
+            {
+                lock(cachesLock)
+                {
+                    IntellisenseCache cache;
+                    if( caches.TryGetValue(summary, out cache) )
+                    {
+                        cache.ReferenceCount--;
+
+                        // Remove unused caches
+                        if( cache.ReferenceCount == 0 )
+                        {
+                            caches.Remove(summary);
+                        }
+                    }
+                }
+            });
+        }
+
+
         /// <summary>
         /// Update the cached autocomplete candidate list when the user connects to a database
-        /// TODO: Update with refactoring/async
         /// </summary>
         /// <param name="details"></param>
         public async Task UpdateAutoCompleteCache(ConnectionDetails details)
         {
             IntellisenseCache cache;
-            if(!caches.TryGetValue(details, out cache))
+            lock(cachesLock)
             {
-                cache = new IntellisenseCache(ConnectionFactory, details);
-                caches[cache.DatabaseInfo] = cache;
+                if(!caches.TryGetValue(details, out cache))
+                {
+                    cache = new IntellisenseCache(ConnectionFactory, details);
+                    caches[cache.DatabaseInfo] = cache;
+                }
+                cache.ReferenceCount++;
             }
             
             await cache.UpdateCache();

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs

@@ -105,9 +105,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             // Register the file open update handler
             WorkspaceService<SqlToolsSettings>.Instance.RegisterTextDocOpenCallback(HandleDidOpenTextDocumentNotification); 
 
-            // register an OnConnection callback 
-            ConnectionService.Instance.RegisterOnConnectionTask(OnConnection);      
-
             // Store the SqlToolsContext for future use
             Context = context;
         }
@@ -305,16 +302,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             CurrentSettings.ScriptAnalysis.Update(newSettings.ScriptAnalysis, CurrentWorkspace.WorkspacePath);
         }
         
-        /// <summary> 
-        /// Callback for when a user connection is done processing 
-        /// </summary> 
-        /// <param name="sqlConnection"></param> 
-        public async Task OnConnection(ConnectionInfo connectionInfo) 
-        { 
-            // TODO consider whether this is needed at all - currently AutoComplete service handles its own updating
-            await Task.FromResult(true); 
-        } 
-        
         #endregion
 
         #region Private Helpers

test/Microsoft.SqlTools.ServiceLayer.Test/Connection/ConnectionServiceTests.cs

@@ -19,6 +19,53 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Connection
     /// </summary>
     public class ConnectionServiceTests
     {
+        /// <summary>
+        /// Verify that when a connection is started for a URI with an already existing
+        /// connection, we disconnect first before connecting.
+        /// </summary>
+        [Fact]
+        public void ConnectingWhenConnectionExistCausesDisconnectThenConnect()
+        {
+            bool callbackInvoked = false;
+
+            // first connect
+            string ownerUri = "file://my/sample/file.sql";
+            var connectionService = TestObjects.GetTestConnectionService();
+            var connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+
+            // verify that we are connected
+            Assert.NotEmpty(connectionResult.ConnectionId);
+
+            // register disconnect callback
+            connectionService.RegisterOnDisconnectTask(
+                (result) => { 
+                    callbackInvoked = true;
+                    return Task.FromResult(true);
+                }
+            );
+
+            // send annother connect request
+            connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+
+            // verify that the event was fired (we disconnected first before connecting)
+            Assert.True(callbackInvoked);
+
+            // verify that we connected again
+            Assert.NotEmpty(connectionResult.ConnectionId);
+        }
+
         /// <summary>
         /// Verify that when connecting with invalid credentials, an error is thrown.
         /// </summary>
@@ -117,6 +164,151 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Connection
             Assert.NotEmpty(connectionResult.ConnectionId);
         }
 
+        /// <summary>
+        /// Verify that we can disconnect from an active connection succesfully
+        /// </summary>
+        [Fact]
+        public void DisconnectFromDatabaseTest()
+        {
+            // first connect
+            string ownerUri = "file://my/sample/file.sql";
+            var connectionService = TestObjects.GetTestConnectionService();
+            var connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+
+            // verify that we are connected
+            Assert.NotEmpty(connectionResult.ConnectionId);
+
+            // send disconnect request
+            var disconnectResult =
+                connectionService
+                .Disconnect(new DisconnectParams()
+                {
+                    OwnerUri = ownerUri
+                });
+            Assert.True(disconnectResult);
+        }
+
+        /// <summary>
+        /// Test that when a disconnect is performed, the callback event is fired
+        /// </summary>
+        [Fact]
+        public void DisconnectFiresCallbackEvent()
+        {
+            bool callbackInvoked = false;
+
+            // first connect
+            string ownerUri = "file://my/sample/file.sql";
+            var connectionService = TestObjects.GetTestConnectionService();
+            var connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+
+            // verify that we are connected
+            Assert.NotEmpty(connectionResult.ConnectionId);
+
+            // register disconnect callback
+            connectionService.RegisterOnDisconnectTask(
+                (result) => { 
+                    callbackInvoked = true;
+                    return Task.FromResult(true);
+                }
+            );
+
+            // send disconnect request
+            var disconnectResult =
+                connectionService
+                .Disconnect(new DisconnectParams()
+                {
+                    OwnerUri = ownerUri
+                });
+            Assert.True(disconnectResult);
+
+            // verify that the event was fired
+            Assert.True(callbackInvoked);
+        }
+
+        /// <summary>
+        /// Test that disconnecting an active connection removes the Owner URI -> ConnectionInfo mapping
+        /// </summary>
+        [Fact]
+        public void DisconnectRemovesOwnerMapping()
+        {
+            // first connect
+            string ownerUri = "file://my/sample/file.sql";
+            var connectionService = TestObjects.GetTestConnectionService();
+            var connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+
+            // verify that we are connected
+            Assert.NotEmpty(connectionResult.ConnectionId);
+
+            // check that the owner mapping exists
+            ConnectionInfo info;
+            Assert.True(connectionService.TryFindConnection(ownerUri, out info));
+
+            // send disconnect request
+            var disconnectResult =
+                connectionService
+                .Disconnect(new DisconnectParams()
+                {
+                    OwnerUri = ownerUri
+                });
+            Assert.True(disconnectResult);
+
+            // check that the owner mapping no longer exists
+            Assert.False(connectionService.TryFindConnection(ownerUri, out info));
+        }
+
+        /// <summary>
+        /// Test that disconnecting validates parameters and doesn't succeed when they are invalid
+        /// </summary>
+        [Theory]
+        [InlineDataAttribute(null)]
+        [InlineDataAttribute("")]
+
+        public void DisconnectValidatesParameters(string disconnectUri)
+        {
+            // first connect
+            string ownerUri = "file://my/sample/file.sql";
+            var connectionService = TestObjects.GetTestConnectionService();
+            var connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+
+            // verify that we are connected
+            Assert.NotEmpty(connectionResult.ConnectionId);
+
+            // send disconnect request
+            var disconnectResult =
+                connectionService
+                .Disconnect(new DisconnectParams()
+                {
+                    OwnerUri = disconnectUri
+                });
+
+            // verify that disconnect failed
+            Assert.False(disconnectResult);
+        }
+
         /// <summary>
         /// Verify that the SQL parser correctly detects errors in text
         /// </summary>
@@ -151,24 +343,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Connection
             var service = TestObjects.GetTestConnectionService();
             var connectParams = TestObjects.GetTestConnectionParams();
 
-            //var endpoint = new Mock<IProtocolEndpoint>();
-            //Func<ConnectParams, RequestContext<ConnectResponse>, Task> connectRequestHandler = null;
-            //endpoint.Setup(e => e.SetRequestHandler(ConnectionRequest.Type, It.IsAny<Func<ConnectParams, RequestContext<ConnectResponse>, Task>>()))
-            //    .Callback<Func<ConnectParams, RequestContext<ConnectResponse>, Task>>(handler => connectRequestHandler = handler);
-
-            // when I initialize the service
-            //service.InitializeService(endpoint.Object);
-
-            // then I expect the handler to be captured
-            //Assert.NotNull(connectRequestHandler);
-
-            // when I call the service
-            //var requestContext = new Mock<RequestContext<ConnectResponse>>();
-
-            //connectRequestHandler(connectParams, requestContext.Object);
-            // then I should get a live connection
-
-            // and then I should have 
             // connect to a database instance 
             var connectionResult = service.Connect(connectParams);
 

test/Microsoft.SqlTools.ServiceLayer.Test/Utility/TestObjects.cs

@@ -322,7 +322,7 @@ namespace Microsoft.SqlTools.Test.Utility
 
         public override void Close()
         {
-            throw new NotImplementedException();
+            // No Op
         }
 
         public override void Open()