From 7f6b357eb00dac65fc6003c564b05518ba38e3ce Mon Sep 17 00:00:00 2001 From: Alan Ren Date: Mon, 19 Jun 2023 16:31:47 -0700 Subject: [PATCH] fix delay in query execution (#2109) * fix delay in query execution * fix test case --- .../Connection/ConnectionService.cs | 36 ++++++++++++------- .../QueryExecution/Query.cs | 2 +- .../Connection/ConnectionServiceTests.cs | 6 ++-- 3 files changed, 28 insertions(+), 16 deletions(-) diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs index edac0010..a999d070 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs @@ -491,14 +491,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection { if (connectionParams != null && connectionParams.Type == ConnectionType.Default && !string.IsNullOrWhiteSpace(connectionParams.OwnerUri)) { - if (connectionParams.OwnerUri.ToLowerInvariant().StartsWith("dashboard://")) + var uri = connectionParams.OwnerUri.ToLowerInvariant(); + if (uri.StartsWith("dashboard://")) { connectionParams.Purpose = ConnectionType.Dashboard; } - else if (connectionParams.OwnerUri.ToLowerInvariant().StartsWith("connection://")) + else if (uri.StartsWith("connection://")) { connectionParams.Purpose = ConnectionType.GeneralConnection; } + else if (uri.StartsWith("untitled:sqlquery") || (uri.StartsWith("file://") && uri.EndsWith(".sql"))) + { + connectionParams.Purpose = ConnectionType.Query; + } } else if (connectionParams != null) { @@ -1044,27 +1049,33 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection { try { - bool disconnect = false; - if (connection.ConnectionString != null){ + bool disconnect = false; + if (connection.ConnectionString != null) + { int totalCount = 0; foreach (KeyValuePair entry in OwnerToConnectionMap) { - foreach (DbConnection value in entry.Value.AllConnections) { - if(value.ConnectionString == connection.ConnectionString) { + foreach (DbConnection value in entry.Value.AllConnections) + { + if (value.ConnectionString == connection.ConnectionString) + { totalCount++; } } } - if(totalCount == 1) { - disconnect = true; + if (totalCount == 1) + { + disconnect = true; } } - else { + else + { disconnect = true; } - - if(disconnect) { + + if (disconnect) + { connection.Close(); } } @@ -1870,7 +1881,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection connInfo.ConnectionDetails.PersistSecurityInfo = true; // turn off connection pool to avoid hold locks on server resources after calling SqlConnection Close method - if (shouldForceDisablePooling) { + if (shouldForceDisablePooling) + { connInfo.ConnectionDetails.Pooling = false; } connInfo.ConnectionDetails.ApplicationName = GetApplicationNameWithFeature(connInfo.ConnectionDetails.ApplicationName, featureName); diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs index a15e5177..9f404baa 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs @@ -448,7 +448,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution } // Locate and setup the connection - queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Query); + queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Default); onErrorAction = OnErrorAction.Ignore; sqlConn = queryConnection as ReliableSqlConnection; if (sqlConn != null) diff --git a/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Connection/ConnectionServiceTests.cs b/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Connection/ConnectionServiceTests.cs index 866eaa3e..5504f8fe 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Connection/ConnectionServiceTests.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Connection/ConnectionServiceTests.cs @@ -45,8 +45,8 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Connection query.Execute(); query.ExecutionTask.Wait(); - // We should see two DbConnections - Assert.AreEqual(2, connectionInfo.CountConnections); + // We should see 1 DbConnections + Assert.AreEqual(1, connectionInfo.CountConnections); // If we run another query query = new Query(Constants.StandardQuery, connectionInfo, new QueryExecutionSettings(), fileStreamFactory); @@ -54,7 +54,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Connection query.ExecutionTask.Wait(); // We should still have 2 DbConnections - Assert.AreEqual(2, connectionInfo.CountConnections); + Assert.AreEqual(1, 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