Simplify connection cleanup and clear pools to prevent lingering connections (#2199)

src/Microsoft.SqlTools.ServiceLayer/Management/Common/DataContainer.cs

@@ -1115,6 +1115,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Management
                     }
                     this.managedConnection = null;
                 }
+                if (this.serverConnection != null)
+                {
+                    if (disposing && this.ownConnection)
+                    {
+                        this.serverConnection.Disconnect();
+                    }
+                    this.serverConnection = null;
+                }
             }
             catch (Exception)
             {

src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/Contracts/DetachDatabaseRequest.cs

@@ -16,6 +16,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement.Contracts
         /// </summary>
         public string ObjectUrn { get; set; }
         /// <summary>
+        /// The target database name.
+        /// </summary>
+        public string Database { get; set; }
+        /// <summary>
         /// URI of the underlying connection for this request
         /// </summary>
         public string ConnectionUri { get; set; }

src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/Contracts/DropDatabaseRequest.cs

@@ -16,6 +16,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement.Contracts
         /// </summary>
         public string ObjectUrn { get; set; }
         /// <summary>
+        /// The target database name.
+        /// </summary>
+        public string Database { get; set; }
+        /// <summary>
         /// URI of the underlying connection for this request
         /// </summary>
         public string ConnectionUri { get; set; }

src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectTypes/Database/DatabaseHandler.cs

@@ -23,6 +23,7 @@ using Microsoft.SqlTools.ServiceLayer.Utility.SqlScriptFormatters;
 using System.Collections.Specialized;
 using Microsoft.SqlTools.SqlCore.Utility;
 using System.Collections.Concurrent;
+using Microsoft.Data.SqlClient;
 
 namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
 {
@@ -149,10 +150,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                 {
                     throw new InvalidOperationException(serverNotExistsError);
                 }
-                try
+
-                {
                 using (var taskHelper = new DatabaseTaskHelper(dataContainer))
-                    using (var context = new DatabaseViewContext(requestParams))
                 {
                     var prototype = taskHelper.Prototype;
                     var azurePrototype = prototype as DatabasePrototypeAzure;
@@ -297,14 +296,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                         databaseViewInfo.LoginNames = new OptionsCollection() { Options = logins.ToArray(), DefaultValueIndex = 0 };
                     }
 
+                    var context = new DatabaseViewContext(requestParams);
                     return Task.FromResult(new InitializeViewResult { ViewInfo = databaseViewInfo, Context = context });
                 }
             }
-                finally
-                {
-                    dataContainer.ServerConnection.Disconnect();
-                }
-            }
         }
 
         public override Task Save(DatabaseViewContext context, DatabaseInfo obj)
@@ -334,10 +329,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
         public string Detach(DetachDatabaseRequestParams detachParams)
         {
             var sqlScript = string.Empty;
-            ConnectionInfo connectionInfo = this.GetConnectionInfo(detachParams.ConnectionUri);
+            using (var dataContainer = CreateDatabaseDataContainer(detachParams.ConnectionUri, detachParams.ObjectUrn, false, detachParams.Database))
-            using (var dataContainer = CreateDatabaseDataContainer(detachParams.ConnectionUri, detachParams.ObjectUrn, false, null))
-            {
-                try
             {
                 var smoDatabase = dataContainer.SqlDialogSubject as Database;
                 if (smoDatabase != null)
@@ -360,6 +352,13 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                                 smoDatabase.DatabaseOptions.UserAccess = SqlServer.Management.Smo.DatabaseUserAccess.Single;
                                 smoDatabase.Alter(TerminationClause.RollbackTransactionsImmediately);
                             }
+                            else
+                            {
+                                // Clear any other connections in the same pool to prevent Detach from hanging
+                                // due to the database still being in use
+                                var sqlConn = dataContainer.ServerConnection.SqlConnectionObject;
+                                SqlConnection.ClearPool(sqlConn);
+                            }
                             smoDatabase.Parent.DetachDatabase(smoDatabase.Name, detachParams.UpdateStatistics);
                         }
                         catch (SmoException)
@@ -380,11 +379,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                     throw new InvalidOperationException($"Provided URN '{detachParams.ObjectUrn}' did not correspond to an existing database.");
                 }
             }
-                finally
-                {
-                    dataContainer.ServerConnection.Disconnect();
-                }
-            }
             return sqlScript;
         }
 
@@ -457,7 +451,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                     {
                         server.ConnectionContext.SqlExecutionModes = originalExecuteMode;
                     }
-                    dataContainer.ServerConnection.Disconnect();
                 }
             }
             return sqlScript;
@@ -470,10 +463,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
         public string Drop(DropDatabaseRequestParams dropParams)
         {
             var sqlScript = string.Empty;
-            ConnectionInfo connectionInfo = this.GetConnectionInfo(dropParams.ConnectionUri);
+            using (var dataContainer = CreateDatabaseDataContainer(dropParams.ConnectionUri, dropParams.ObjectUrn, false, dropParams.Database))
-            using (var dataContainer = CreateDatabaseDataContainer(dropParams.ConnectionUri, dropParams.ObjectUrn, false, null))
-            {
-                try
             {
                 var smoDatabase = dataContainer.SqlDialogSubject as Database;
                 if (smoDatabase != null)
@@ -498,6 +488,13 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                             smoDatabase.DatabaseOptions.UserAccess = SqlServer.Management.Smo.DatabaseUserAccess.Single;
                             smoDatabase.Alter(TerminationClause.RollbackTransactionsImmediately);
                         }
+                        else
+                        {
+                            // Clear any other connections in the same pool to prevent Drop from hanging
+                            // due to the database still being in use
+                            var sqlConn = dataContainer.ServerConnection.SqlConnectionObject;
+                            SqlConnection.ClearPool(sqlConn);
+                        }
                         if (dropParams.DeleteBackupHistory)
                         {
                             server.DeleteBackupHistory(smoDatabase.Name);
@@ -541,11 +538,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                     throw new InvalidOperationException($"Provided URN '{dropParams.ObjectUrn}' did not correspond to an existing database.");
                 }
             }
-                finally
-                {
-                    dataContainer.ServerConnection.Disconnect();
-                }
-            }
             return sqlScript;
         }
 
@@ -590,8 +582,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                 {
                     throw new InvalidOperationException(serverNotExistsError);
                 }
-                try
-                {
                 using (var taskHelper = new DatabaseTaskHelper(dataContainer))
                 {
                     DatabasePrototype prototype = taskHelper.Prototype;
@@ -747,7 +737,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                         }
 
                         // Remove the file
-                            foreach (KeyValuePair<int,DatabaseFilePrototype> removedfile in fileDict)
+                        foreach (KeyValuePair<int, DatabaseFilePrototype> removedfile in fileDict)
                         {
                             removedfile.Value.Removed = true;
                         }
@@ -850,11 +840,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                     return sqlScript;
                 }
             }
-                finally
-                {
-                    dataContainer.ServerConnection.Disconnect();
-                }
-            }
         }
 
         private Autogrowth GetAutogrowth(DatabasePrototype prototype, DatabaseFile file)

src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectTypes/Server/ServerHandler.cs

@@ -32,9 +32,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
         public override Task<InitializeViewResult> InitializeObjectView(InitializeViewRequestParams requestParams)
         {
             ConnectionInfo connInfo = this.GetConnectionInfo(requestParams.ConnectionUri);
-
-            using (var context = new ServerViewContext(requestParams, ConnectionService.OpenServerConnection(connInfo, ObjectManagementService.ApplicationName)))
-            {
             CDataContainer dataContainer = CDataContainer.CreateDataContainer(connInfo, databaseExists: true);
 
             ServerPrototype prototype = new ServerPrototype(dataContainer);
@@ -68,9 +65,9 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                     NumaNodes = prototype.NumaNodes
                 };
             }
+            var context = new ServerViewContext(requestParams);
             return Task.FromResult(new InitializeViewResult { ViewInfo = this.serverViewInfo, Context = context });
         }
-        }
 
         public override Task Save(ServerViewContext context, ServerInfo obj)
         {

src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectTypes/Server/ServerViewContext.cs

@@ -3,16 +3,12 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
-using Microsoft.SqlServer.Management.Common;
-
 namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
 {
     public class ServerViewContext : SqlObjectViewContext
     {
-        public ServerConnection Connection { get; }
+        public ServerViewContext(Contracts.InitializeViewRequestParams parameters) : base(parameters)
-        public ServerViewContext(Contracts.InitializeViewRequestParams parameters, ServerConnection connection) : base(parameters)
         {
-            this.Connection = connection;
         }
 
         public override void Dispose()