Fixing the bug with connections on database make restore fail (#473)

* closing the connections that don't need to be open and keeping track of the connections that should stay open

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -18,8 +18,6 @@ using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
-using Microsoft.SqlTools.ServiceLayer.SqlContext;
-using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlTools.Utility;
 
@@ -53,7 +51,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// The SQL connection factory object
         /// </summary>
         private ISqlConnectionFactory connectionFactory;
-           
+
+        private DatabaseLocksManager lockedDatabaseManager;
+
         private readonly Dictionary<string, ConnectionInfo> ownerToConnectionMap = new Dictionary<string, ConnectionInfo>();
 
         /// <summary>
@@ -65,7 +65,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
 
         private readonly object cancellationTokenSourceLock = new object();
 
-        private ConnectedBindingQueue connectionQueue = new ConnectedBindingQueue(needsMetadata: false);
+        private ConcurrentDictionary<string, IConnectedBindingQueue> connectedQueues = new ConcurrentDictionary<string, IConnectedBindingQueue>();
 
         /// <summary>
         /// Map from script URIs to ConnectionInfo objects
@@ -79,6 +79,25 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             }
         }
 
+        /// <summary>
+        /// Database Lock manager instance
+        /// </summary>
+        internal DatabaseLocksManager LockedDatabaseManager
+        {
+            get
+            {
+                if (lockedDatabaseManager == null)
+                {
+                    lockedDatabaseManager = DatabaseLocksManager.Instance;
+                }
+                return lockedDatabaseManager;
+            }
+            set
+            {
+                this.lockedDatabaseManager = value;
+            }
+        }
+
         /// <summary>
         /// Service host object for sending/receiving requests/events.
         /// Internal for testing purposes.
@@ -92,20 +111,63 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <summary>
         /// Gets the connection queue
         /// </summary>
-        internal ConnectedBindingQueue ConnectionQueue
+        internal IConnectedBindingQueue ConnectionQueue
         {
             get
             {
-                return this.connectionQueue;
+                return this.GetConnectedQueue("Default");
             }
         }
 
+
         /// <summary>
         /// Default constructor should be private since it's a singleton class, but we need a constructor
         /// for use in unit test mocking.
         /// </summary>
         public ConnectionService()
         {
+            var defaultQueue = new ConnectedBindingQueue(needsMetadata: false);
+            connectedQueues.AddOrUpdate("Default", defaultQueue, (key, old) => defaultQueue);
+            this.LockedDatabaseManager.ConnectionService = this;
+        }
+
+        /// <summary>
+        /// Returns a connection queue for given type
+        /// </summary>
+        /// <param name="type"></param>
+        /// <returns></returns>
+        public IConnectedBindingQueue GetConnectedQueue(string type)
+        {
+            IConnectedBindingQueue connectedBindingQueue;
+            if (connectedQueues.TryGetValue(type, out connectedBindingQueue))
+            {
+                return connectedBindingQueue;
+            }
+            return null;
+        }
+
+        /// <summary>
+        /// Returns all the connection queues
+        /// </summary>
+        public IEnumerable<IConnectedBindingQueue> ConnectedQueues
+        {
+            get
+            {
+                return this.connectedQueues.Values;
+            }
+        }
+
+        /// <summary>
+        /// Register a new connection queue if not already registered
+        /// </summary>
+        /// <param name="type"></param>
+        /// <param name="connectedQueue"></param>
+        public virtual void RegisterConnectedQueue(string type, IConnectedBindingQueue connectedQueue)
+        {
+            if (!connectedQueues.ContainsKey(type))
+            {
+                connectedQueues.AddOrUpdate(type, connectedQueue, (key, old) => connectedQueue);
+            }
         }
 
         /// <summary>
@@ -243,6 +305,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             // Invoke callback notifications          
             InvokeOnConnectionActivities(connectionInfo, connectionParams);
 
+            if(connectionParams.Type == ConnectionType.ObjectExplorer)
+            {
+                DbConnection connection;
+                if (connectionInfo.TryGetConnection(ConnectionType.ObjectExplorer, out connection))
+                {
+                    // OE doesn't need to keep the connection open
+                    connection.Close();
+                }
+            }
             return completeParams;
         }
 
@@ -359,9 +430,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             DbConnection connection = null;
             CancelTokenKey cancelKey = new CancelTokenKey { OwnerUri = connectionParams.OwnerUri, Type = connectionParams.Type };
             ConnectionCompleteParams response = new ConnectionCompleteParams { OwnerUri = connectionInfo.OwnerUri, Type = connectionParams.Type };
+            bool? currentPooling = connectionInfo.ConnectionDetails.Pooling;
 
             try
             {
+                connectionInfo.ConnectionDetails.Pooling = false;
                 // build the connection string from the input parameters
                 string connectionString = BuildConnectionString(connectionInfo.ConnectionDetails);
 
@@ -382,7 +455,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                     }
                     cancelTupleToCancellationTokenSourceMap[cancelKey] = source;
                 }
-                    
+                
                 // Open the connection
                 await connection.OpenAsync(source.Token);
             }
@@ -419,6 +492,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                     }
                     source?.Dispose();
                 }
+                if (connectionInfo != null && connectionInfo.ConnectionDetails != null)
+                {
+                    connectionInfo.ConnectionDetails.Pooling = currentPooling;
+                }
             }
 
             // Return null upon success
@@ -1158,7 +1235,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 // turn off connection pool to avoid hold locks on server resources after calling SqlConnection Close method
                 connInfo.ConnectionDetails.Pooling = false;
 
-                // generate connection string        
+                // generate connection string
                 string connectionString = ConnectionService.BuildConnectionString(connInfo.ConnectionDetails);
 
                 // restore original values
@@ -1167,7 +1244,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 connInfo.ConnectionDetails.Pooling = originalPooling;
 
                 // open a dedicated binding server connection
-                SqlConnection sqlConn = new SqlConnection(connectionString); 
+                SqlConnection sqlConn = new SqlConnection(connectionString);
                 sqlConn.Open();
                 return sqlConn;
             }