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.IntegrationTests/Connection/ConnectionServiceTests.cs

@@ -0,0 +1,104 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.Collections.Generic;
+using System.Data.Common;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
+using Microsoft.SqlTools.Test.Utility;
+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.IntegrationTests.Connection
+{
+    /// <summary>
+    /// Tests for the ServiceHost Connection Service tests that require a live database connection
+    /// </summary>
+    public class ConnectionServiceTests
+    {
+        [Fact]
+        public async Task RunningMultipleQueriesCreatesOnlyOneConnection()
+        {
+            // Connect/disconnect twice to ensure reconnection can occur
+            ConnectionService service = ConnectionService.Instance;
+            service.OwnerToConnectionMap.Clear();
+            for (int i = 0; i < 2; i++)
+            {
+                var result = await TestObjects.InitLiveConnectionInfo();
+                ConnectionInfo connectionInfo = result.ConnectionInfo;
+                string uri = connectionInfo.OwnerUri;
+
+                // We should see one ConnectionInfo and one DbConnection
+                Assert.Equal(1, connectionInfo.CountConnections);
+                Assert.Equal(1, service.OwnerToConnectionMap.Count);
+
+                // If we run a query
+                var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+                Query query = new Query(Common.StandardQuery, connectionInfo, new QueryExecutionSettings(), fileStreamFactory);
+                query.Execute();
+                query.ExecutionTask.Wait();
+
+                // We should see two DbConnections
+                Assert.Equal(2, connectionInfo.CountConnections);
+
+                // If we run another query
+                query = new Query(Common.StandardQuery, connectionInfo, new QueryExecutionSettings(), fileStreamFactory);
+                query.Execute();
+                query.ExecutionTask.Wait();
+
+                // We should still have 2 DbConnections
+                Assert.Equal(2, connectionInfo.CountConnections);
+
+                // If we disconnect, we should remain in a consistent state to do it over again
+                // e.g. loop and do it over again
+                service.Disconnect(new DisconnectParams() { OwnerUri = connectionInfo.OwnerUri });
+
+                // We should be left with an empty connection map
+                Assert.Equal(0, service.OwnerToConnectionMap.Count);
+            }
+        }
+
+        [Fact]
+        public async Task DatabaseChangesAffectAllConnections()
+        {
+            // If we make a connection to a live database 
+            ConnectionService service = ConnectionService.Instance;
+            var result = await TestObjects.InitLiveConnectionInfo();
+            ConnectionInfo connectionInfo = result.ConnectionInfo;
+            ConnectionDetails details = connectionInfo.ConnectionDetails;
+            string uri = connectionInfo.OwnerUri;
+            string initialDatabaseName = details.DatabaseName;
+            string newDatabaseName = "tempdb";
+            string changeDatabaseQuery = "use " + newDatabaseName;
+
+            // Then run any query to create a query DbConnection
+            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+            Query query = new Query(Common.StandardQuery, connectionInfo, new QueryExecutionSettings(), fileStreamFactory);
+            query.Execute();
+            query.ExecutionTask.Wait();
+
+            // All open DbConnections (Query and Default) should have initialDatabaseName as their database
+            foreach (DbConnection connection in connectionInfo.AllConnections)
+            {
+                Assert.Equal(connection.Database, initialDatabaseName);
+            }
+
+            // If we run a query to change the database
+            query = new Query(changeDatabaseQuery, connectionInfo, new QueryExecutionSettings(), fileStreamFactory);
+            query.Execute();
+            query.ExecutionTask.Wait();
+
+            // All open DbConnections (Query and Default) should have newDatabaseName as their database
+            foreach (DbConnection connection in connectionInfo.AllConnections)
+            {
+                Assert.Equal(connection.Database, newDatabaseName);
+            }
+        }
+    }
+}

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Connection/ReliableConnectionTests.cs

@@ -685,19 +685,21 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Connection
         {
             var result = await TestObjects.InitLiveConnectionInfo();
             ConnectionInfo connInfo = result.ConnectionInfo;
+            DbConnection connection = connInfo.ConnectionTypeToConnectionMap[ConnectionType.Default];
 
-            Assert.True(ReliableConnectionHelper.IsAuthenticatingDatabaseMaster(connInfo.SqlConnection));
+
+            Assert.True(ReliableConnectionHelper.IsAuthenticatingDatabaseMaster(connection));
 
             SqlConnectionStringBuilder builder = new SqlConnectionStringBuilder();
             Assert.True(ReliableConnectionHelper.IsAuthenticatingDatabaseMaster(builder));
             ReliableConnectionHelper.TryAddAlwaysOnConnectionProperties(builder, new SqlConnectionStringBuilder());
 
-            Assert.NotNull(ReliableConnectionHelper.GetServerName(connInfo.SqlConnection));
-            Assert.NotNull(ReliableConnectionHelper.ReadServerVersion(connInfo.SqlConnection));
+            Assert.NotNull(ReliableConnectionHelper.GetServerName(connection));
+            Assert.NotNull(ReliableConnectionHelper.ReadServerVersion(connection));
             
-            Assert.NotNull(ReliableConnectionHelper.GetAsSqlConnection(connInfo.SqlConnection));
+            Assert.NotNull(ReliableConnectionHelper.GetAsSqlConnection(connection));
 
-            ReliableConnectionHelper.ServerInfo info = ReliableConnectionHelper.GetServerVersion(connInfo.SqlConnection);
+            ReliableConnectionHelper.ServerInfo info = ReliableConnectionHelper.GetServerVersion(connection);
             Assert.NotNull(ReliableConnectionHelper.IsVersionGreaterThan2012RTM(info));
         }
 
@@ -728,8 +730,10 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Connection
         {
             var result = await TestObjects.InitLiveConnectionInfo();
             ConnectionInfo connInfo = result.ConnectionInfo;
+            DbConnection dbConnection;
+            connInfo.TryGetConnection(ConnectionType.Default, out dbConnection);
 
-            var connection = connInfo.SqlConnection as ReliableSqlConnection;
+            var connection = dbConnection as ReliableSqlConnection;
             var command = new ReliableSqlConnection.ReliableSqlCommand(connection);
             Assert.NotNull(command.Connection);
 

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/QueryExecution/DataStorage/StorageDataReaderTests.cs

@@ -19,8 +19,10 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.QueryExecution.DataSt
         private async Task<StorageDataReader> GetTestStorageDataReader(string query)
         {
             var result = await TestObjects.InitLiveConnectionInfo();
+            DbConnection connection;
+            result.ConnectionInfo.TryGetConnection(ConnectionType.Default, out connection);
 
-            var command = result.ConnectionInfo.SqlConnection.CreateCommand();
+            var command = connection.CreateCommand();
             command.CommandText = query;
             DbDataReader reader = command.ExecuteReader();
 

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

@@ -0,0 +1,112 @@
+using System;
+using System.Collections.Generic;
+using System.Data.Common;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Test.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
+using Microsoft.SqlTools.Test.Utility;
+using Xunit;
+
+namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.QueryExecution
+{
+    public class ExecuteTests
+    {
+        [Fact]
+        public async Task RollbackTransactionFailsWithoutBeginTransaction()
+        {
+            const string refactorText = "ROLLBACK TRANSACTION";
+
+            // Given a connection to a live database
+            var result = await TestObjects.InitLiveConnectionInfo();
+            ConnectionInfo connInfo = result.ConnectionInfo;
+            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+
+            // If I run a "ROLLBACK TRANSACTION" query
+            Query query = new Query(refactorText, connInfo, new QueryExecutionSettings(), fileStreamFactory);
+            query.Execute();
+            query.ExecutionTask.Wait();
+
+            // There should be an error
+            Assert.True(query.Batches[0].HasError);
+        }
+
+        [Fact]
+        public async Task TransactionsSucceedAcrossQueries()
+        {
+            const string beginText = "BEGIN TRANSACTION";
+            const string rollbackText = "ROLLBACK TRANSACTION";
+
+            // Given a connection to a live database
+            var result = await TestObjects.InitLiveConnectionInfo();
+            ConnectionInfo connInfo = result.ConnectionInfo;
+            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+
+            // If I run a "BEGIN TRANSACTION" query
+            CreateAndExecuteQuery(beginText, connInfo, fileStreamFactory);
+
+            // Then I run a "ROLLBACK TRANSACTION" query, there should be no errors
+            Query rollbackQuery = CreateAndExecuteQuery(rollbackText, connInfo, fileStreamFactory);
+            Assert.False(rollbackQuery.Batches[0].HasError);
+        }
+
+        [Fact]
+        public async Task TempTablesPersistAcrossQueries()
+        {
+            const string createTempText = "CREATE TABLE #someTempTable (id int)";
+            const string insertTempText = "INSERT INTO #someTempTable VALUES(1)";
+
+            // Given a connection to a live database
+            var result = await TestObjects.InitLiveConnectionInfo();
+            ConnectionInfo connInfo = result.ConnectionInfo;
+            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+
+            // If I run a query creating a temp table
+            CreateAndExecuteQuery(createTempText, connInfo, fileStreamFactory);
+
+            // Then I run a different query using that temp table, there should be no errors
+            Query insertTempQuery = CreateAndExecuteQuery(insertTempText, connInfo, fileStreamFactory);
+            Assert.False(insertTempQuery.Batches[0].HasError);
+        }
+
+        [Fact]
+        public async Task DatabaseChangesWhenCallingUseDatabase()
+        {
+            const string master = "master";
+            const string tempdb = "tempdb";
+            const string useQuery = "USE {0}";
+
+            // Given a connection to a live database
+            var result = await TestObjects.InitLiveConnectionInfo();
+            ConnectionInfo connInfo = result.ConnectionInfo;
+            DbConnection connection;
+            connInfo.TryGetConnection(ConnectionType.Default, out connection);
+
+            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+
+            // If I use master, the current database should be master
+            CreateAndExecuteQuery(string.Format(useQuery, master), connInfo, fileStreamFactory);
+            Assert.Equal(master, connection.Database);
+
+            // If I use tempdb, the current database should be tempdb
+            CreateAndExecuteQuery(string.Format(useQuery, tempdb), connInfo, fileStreamFactory);
+            Assert.Equal(tempdb, connection.Database);
+
+            // If I switch back to master, the current database should be master
+            CreateAndExecuteQuery(string.Format(useQuery, master), connInfo, fileStreamFactory);
+            Assert.Equal(master, connection.Database);
+        }
+
+        public Query CreateAndExecuteQuery(string queryText, ConnectionInfo connectionInfo, IFileStreamFactory fileStreamFactory)
+        {
+            Query query = new Query(queryText, connectionInfo, new QueryExecutionSettings(), fileStreamFactory);
+            query.Execute();
+            query.ExecutionTask.Wait();
+            return query;
+        }
+    }
+}

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,