Update connection logic to handle multiple connections per URI (#176)

* Add CancelTokenKey for uniquely identifying cancelations of Connections associated with an OwnerUri and ConnectionType string.

* Update ConnectionInfo to use ConcurrentDictionary of DbConnection instances. Add wrapper functions for the ConcurrentDictionary. 

* Refactor Connect and Disconnect in ConnectionService.

* Update ConnectionService: Handle multiple connections per ConnectionInfo. Handle cancelation tokens uniquely identified with CancelTokenKey. Add GetOrOpenConnection() for other services to request an existing or create a new DbConnection.

* Add ConnectionType.cs for ConnectionType strings.

* Add ConnectionType string to ConnectParams, ConnectionCompleteNotification, DisconnectParams.

* Update Query ExecuteInternal to use the dedicated query connection and GetOrOpenConnection().

* Update test library to account for multiple connections in ConnectionInfo.

* Write tests ensuring multiple connections don’t create redundant data.

* Write tests ensuring database changes affect all connections of a given ConnectionInfo.

* Write tests for TRANSACTION statements and temp tables.

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

@@ -19,6 +19,10 @@ using Microsoft.SqlTools.Test.Utility;
 using Moq;
 using Moq.Protected;
 using Xunit;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Test.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.Test.Connection
 {
@@ -994,5 +998,161 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Connection
             Assert.NotNull(errorMessage);
             Assert.NotEmpty(errorMessage);
         }
+
+        [Fact]
+        public async void ConnectingTwiceWithTheSameUriDoesNotCreateAnotherDbConnection()
+        {
+            // Setup the connect and disconnect params
+            var connectParamsSame1 = new ConnectParams()
+            {
+                OwnerUri = "connectParamsSame",
+                Connection = TestObjects.GetTestConnectionDetails()
+            };
+            var connectParamsSame2 = new ConnectParams()
+            {
+                OwnerUri = "connectParamsSame",
+                Connection = TestObjects.GetTestConnectionDetails()
+            };
+            var disconnectParamsSame = new DisconnectParams()
+            {
+                OwnerUri = connectParamsSame1.OwnerUri
+            };
+            var connectParamsDifferent = new ConnectParams()
+            {
+                OwnerUri = "connectParamsDifferent",
+                Connection = TestObjects.GetTestConnectionDetails()
+            };
+            var disconnectParamsDifferent = new DisconnectParams()
+            {
+                OwnerUri = connectParamsDifferent.OwnerUri
+            };
+
+            // Given a request to connect to a database, there should be no initial connections in the map
+            var service = TestObjects.GetTestConnectionService();
+            Dictionary<string, ConnectionInfo> ownerToConnectionMap = service.OwnerToConnectionMap;
+            Assert.Equal(0, ownerToConnectionMap.Count);
+
+            // If we connect to the service, there should be 1 connection
+            await service.Connect(connectParamsSame1);
+            Assert.Equal(1, ownerToConnectionMap.Count);
+
+            // If we connect again with the same URI, there should still be 1 connection
+            await service.Connect(connectParamsSame2);
+            Assert.Equal(1, ownerToConnectionMap.Count);
+
+            // If we connect with a different URI, there should be 2 connections
+            await service.Connect(connectParamsDifferent);
+            Assert.Equal(2, ownerToConnectionMap.Count);
+
+            // If we disconenct with the unique URI, there should be 1 connection
+            service.Disconnect(disconnectParamsDifferent);
+            Assert.Equal(1, ownerToConnectionMap.Count);
+
+            // If we disconenct with the duplicate URI, there should be 0 connections
+            service.Disconnect(disconnectParamsSame);
+            Assert.Equal(0, ownerToConnectionMap.Count);
+        }
+
+        [Fact]
+        public async void DbConnectionDoesntLeakUponDisconnect()
+        {
+            // If we connect with a single URI and 2 connection types
+            var connectParamsDefault = new ConnectParams()
+            {
+                OwnerUri = "connectParams",
+                Connection = TestObjects.GetTestConnectionDetails(),
+                Type = ConnectionType.Default
+            };
+            var connectParamsQuery = new ConnectParams()
+            {
+                OwnerUri = "connectParams",
+                Connection = TestObjects.GetTestConnectionDetails(),
+                Type = ConnectionType.Query
+            };
+            var disconnectParams = new DisconnectParams()
+            {
+                OwnerUri = connectParamsDefault.OwnerUri
+            };
+            var service = TestObjects.GetTestConnectionService();
+            await service.Connect(connectParamsDefault);
+            await service.Connect(connectParamsQuery);
+
+            // We should have one ConnectionInfo and 2 DbConnections 
+            ConnectionInfo connectionInfo = service.OwnerToConnectionMap[connectParamsDefault.OwnerUri];
+            Assert.Equal(2, connectionInfo.CountConnections);
+            Assert.Equal(1, service.OwnerToConnectionMap.Count);
+
+            // If we record when the Default connecton calls Close()
+            bool defaultDisconnectCalled = false;
+            var mockDefaultConnection = new Mock<DbConnection> { CallBase = true };
+            mockDefaultConnection.Setup(x => x.Close())
+                .Callback(() =>
+                {
+                    defaultDisconnectCalled = true;
+                });
+            connectionInfo.ConnectionTypeToConnectionMap[ConnectionType.Default] = mockDefaultConnection.Object;
+
+            // And when the Query connecton calls Close()
+            bool queryDisconnectCalled = false;
+            var mockQueryConnection = new Mock<DbConnection> { CallBase = true };
+            mockQueryConnection.Setup(x => x.Close())
+                .Callback(() =>
+                {
+                    queryDisconnectCalled = true;
+                });
+            connectionInfo.ConnectionTypeToConnectionMap[ConnectionType.Query] = mockQueryConnection.Object;
+
+            // If we disconnect all open connections with the same URI as used above 
+            service.Disconnect(disconnectParams);
+
+            // Close() should have gotten called for both DbConnections
+            Assert.True(defaultDisconnectCalled);
+            Assert.True(queryDisconnectCalled);
+
+            // And the maps that hold connection data should be empty
+            Assert.Equal(0, connectionInfo.CountConnections);
+            Assert.Equal(0, service.OwnerToConnectionMap.Count);
+        }
+
+        [Fact]
+        public async void ClosingQueryConnectionShouldLeaveDefaultConnectionOpen()
+        {
+            // Setup the connect and disconnect params
+            var connectParamsDefault = new ConnectParams()
+            {
+                OwnerUri = "connectParamsSame",
+                Connection = TestObjects.GetTestConnectionDetails(),
+                Type = ConnectionType.Default
+            };
+            var connectParamsQuery = new ConnectParams()
+            {
+                OwnerUri = connectParamsDefault.OwnerUri,
+                Connection = TestObjects.GetTestConnectionDetails(),
+                Type = ConnectionType.Query
+            };
+            var disconnectParamsQuery = new DisconnectParams()
+            {
+                OwnerUri = connectParamsDefault.OwnerUri,
+                Type = connectParamsQuery.Type
+            };
+
+            // If I connect a Default and a Query connection
+            var service = TestObjects.GetTestConnectionService();
+            Dictionary<string, ConnectionInfo> ownerToConnectionMap = service.OwnerToConnectionMap;
+            await service.Connect(connectParamsDefault);
+            await service.Connect(connectParamsQuery);
+            ConnectionInfo connectionInfo = service.OwnerToConnectionMap[connectParamsDefault.OwnerUri];
+
+            // There should be 2 connections in the map
+            Assert.Equal(2, connectionInfo.CountConnections);
+
+            // If I Disconnect only the Query connection, there should be 1 connection in the map
+            service.Disconnect(disconnectParamsQuery);
+            Assert.Equal(1, connectionInfo.CountConnections);
+
+            // If I reconnect, there should be 2 again
+            await service.Connect(connectParamsQuery);
+            Assert.Equal(2, connectionInfo.CountConnections);
+        }
     }
 }

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs

@@ -82,6 +82,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         public static Query GetBasicExecutedQuery()
         {
             ConnectionInfo ci = CreateTestConnectionInfo(new[] {StandardTestData}, false);
+
+            // Query won't be able to request a new query DbConnection unless the ConnectionService has a 
+            // ConnectionInfo with the same URI as the query, so we will manually set it
+            ConnectionService.Instance.OwnerToConnectionMap[ci.OwnerUri] = ci;
+
             Query query = new Query(StandardQuery, ci, new QueryExecutionSettings(), GetFileStreamFactory(new Dictionary<string, byte[]>()));
             query.Execute();
             query.ExecutionTask.Wait();
@@ -222,6 +227,23 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             return new ConnectionInfo(CreateMockFactory(data, throwOnRead), OwnerUri, StandardConnectionDetails);
         }
 
+        public static ConnectionInfo CreateConnectedConnectionInfo(Dictionary<string, string>[][] data, bool throwOnRead, string type = ConnectionType.Default)
+        {
+            ConnectionService connectionService = ConnectionService.Instance;
+            connectionService.OwnerToConnectionMap.Clear();
+            connectionService.ConnectionFactory = CreateMockFactory(data, throwOnRead);
+
+            ConnectParams connectParams = new ConnectParams()
+            {
+                Connection = StandardConnectionDetails,
+                OwnerUri = Common.OwnerUri,
+                Type = type
+            };
+
+            connectionService.Connect(connectParams).Wait();
+            return connectionService.OwnerToConnectionMap[OwnerUri];
+        }
+
         #endregion
 
         #region Service Mocking
@@ -233,12 +255,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // Create a place for the temp "files" to be written
             storage = new Dictionary<string, byte[]>();
 
-            // Create the connection factory with the dataset
-            var factory = CreateTestConnectionInfo(data, throwOnRead).Factory;
-
             // Mock the connection service
             var connectionService = new Mock<ConnectionService>();
-            ConnectionInfo ci = new ConnectionInfo(factory, OwnerUri, StandardConnectionDetails);
+            ConnectionInfo ci = CreateConnectedConnectionInfo(data, throwOnRead);
             ConnectionInfo outValMock;
             connectionService
                 .Setup(service => service.TryFindConnection(It.IsAny<string>(), out outValMock))

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs

@@ -7,10 +7,16 @@
 
 using System.Data.Common;
 using System.Threading.Tasks;
+using System;
+using System.Collections.Generic;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
+using Microsoft.SqlTools.Test.Utility;
+using Xunit;
 
 namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 {
@@ -24,12 +30,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // If:
             // ... I create a query with a udt column in the result set
             ConnectionInfo connectionInfo = TestObjects.GetTestConnectionInfo();
-            Query query = new Query(Common.UdtQuery, connectionInfo, new QueryExecutionSettings(), Common.GetFileStreamFactory());
+            Query query = new Query(Common.UdtQuery, connectionInfo, new QueryExecutionSettings(), Common.GetFileStreamFactory(new Dictionary<string, byte[]>()));
 
             // If:
             // ... I then execute the query
             DateTime startTime = DateTime.Now;
-            query.Execute().Wait();
+            query.Execute();
+            query.ExecutionTask.Wait();
 
             // Then:
             // ... The query should complete within 2 seconds since retry logic should not kick in

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Execution/QueryTests.cs

@@ -163,7 +163,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
 
             // If:
             // ... I create a query from two batches (with separator)
-            ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
+            ConnectionInfo ci = Common.CreateConnectedConnectionInfo(null, false);
+
             string queryText = string.Format("{0}\r\nGO\r\n{0}", Common.StandardQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(queryText, ci, new QueryExecutionSettings(), fileStreamFactory);
@@ -280,6 +281,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // If:
             // ... I create a query from an invalid batch
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, true);
+            ConnectionService.Instance.OwnerToConnectionMap[ci.OwnerUri] = ci;
+
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(Common.InvalidQuery, ci, new QueryExecutionSettings(), fileStreamFactory);
             BatchCallbackHelper(query,