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>
@@ -127,16 +137,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
@@ -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);
@@ -202,6 +243,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
@@ -226,6 +275,27 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 await requestContext.SendError(ex.Message);
             }
         }
+
+        /// <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, 

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

@@ -7,57 +7,57 @@ using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
 {
-    /// <summary>
-    /// Parameters for the Connect Request.
+    /// <summary>
+    /// Parameters for the Connect Request.
     /// </summary>
-    public class ConnectParams
-    {
-        /// <summary>
-        /// 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;  }
-        /// <summary>
-        /// Contains the required parameters to initialize a connection to a database.
-        /// A connection will identified by its server name, database name and user name.
-        /// This may be changed in the future to support multiple connections with different 
-        /// connection properties to the same database.
-        /// </summary>
-        public ConnectionDetails Connection { get; set; }
+    public class ConnectParams
+    {
+        /// <summary>
+        /// 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;  }
+        /// <summary>
+        /// Contains the required parameters to initialize a connection to a database.
+        /// A connection will identified by its server name, database name and user name.
+        /// This may be changed in the future to support multiple connections with different 
+        /// connection properties to the same database.
+        /// </summary>
+        public ConnectionDetails Connection { get; set; }
     }
 
-    /// <summary>
-    /// Parameters for the Disconnect Request.
+    /// <summary>
+    /// Parameters for the Disconnect Request.
     /// </summary>
-    public class DisconnectParams
-    {
-        /// <summary>
-        /// 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 class DisconnectParams
+    {
+        /// <summary>
+        /// 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; }
     }
 
-    /// <summary>
-    /// Parameters for the ConnectionChanged Notification.
+    /// <summary>
+    /// Parameters for the ConnectionChanged Notification.
     /// </summary>
-    public class ConnectionChangedParams
-    {
-        /// <summary>
-        /// 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; }
-        /// <summary>
-        /// Contains the high-level properties about the connection, for display to the user.
-        /// </summary>
-        public ConnectionSummary Connection { get; set; }
+    public class ConnectionChangedParams
+    {
+        /// <summary>
+        /// 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; }
+        /// <summary>
+        /// Contains the high-level properties about the connection, for display to the user.
+        /// </summary>
+        public ConnectionSummary Connection { get; set; }
     }
 
-    /// <summary>
-    /// Provides high level information about a connection.
+    /// <summary>
+    /// Provides high level information about a connection.
     /// </summary>
-    public class ConnectionSummary
+    public class ConnectionSummary
     {
         /// <summary>
         /// Gets or sets the connection server name
@@ -72,7 +72,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
         /// <summary>
         /// Gets or sets the connection user name
         /// </summary>
-        public string UserName { get; set; }
+        public string UserName { get; set; }
     }
     /// <summary>
     /// Message format for the initial connection request
@@ -102,8 +102,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
         /// Gets or sets any connection error messages
         /// </summary>
         public string Messages { get; set; }
-    }
-
+    }
+
     /// <summary>
     /// Connect request mapping entry 
     /// </summary>
@@ -122,8 +122,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
         public static readonly
             RequestType<DisconnectParams, bool> Type =
             RequestType<DisconnectParams, bool>.Create("connection/disconnect");
-    }
-
+    }
+
     /// <summary>
     /// ConnectionChanged notification mapping entry 
     /// </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

@@ -103,10 +103,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             WorkspaceService<SqlToolsSettings>.Instance.RegisterTextDocChangeCallback(HandleDidChangeTextDocumentNotification);
 
             // Register the file open update handler
-            WorkspaceService<SqlToolsSettings>.Instance.RegisterTextDocOpenCallback(HandleDidOpenTextDocumentNotification);
-
-            // register an OnConnection callback 
-            ConnectionService.Instance.RegisterOnConnectionTask(OnConnection);      
+            WorkspaceService<SqlToolsSettings>.Instance.RegisterTextDocOpenCallback(HandleDidOpenTextDocumentNotification); 
 
             // 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()