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;
             }

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

@@ -16,5 +16,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         public const string Default = "Default";
         public const string Query = "Query";
         public const string Edit = "Edit";
+        public const string ObjectExplorer = "ObjectExplorer";
     }
 }

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

@@ -0,0 +1,28 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+
+namespace Microsoft.SqlTools.ServiceLayer.Connection
+{
+    public class DatabaseFullAccessException: Exception
+    {
+        public DatabaseFullAccessException()
+            : base()
+        {
+        }
+
+        public DatabaseFullAccessException(string message, Exception exception)
+            : base(message, exception)
+        {
+        }
+
+
+        public DatabaseFullAccessException(string message)
+            : base(message)
+        {
+        }
+    }
+}

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

@@ -0,0 +1,113 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.ServiceLayer.LanguageServices;
+using System;
+using System.Collections.Generic;
+using System.Threading;
+
+namespace Microsoft.SqlTools.ServiceLayer.Connection
+{
+    public class DatabaseLocksManager: IDisposable
+    {
+        internal DatabaseLocksManager(int waitToGetFullAccess)
+        {
+            this.waitToGetFullAccess = waitToGetFullAccess;
+        }
+
+        private static DatabaseLocksManager instance = new DatabaseLocksManager(DefaultWaitToGetFullAccess);
+
+        public static DatabaseLocksManager Instance
+        {
+            get
+            {
+                return instance;
+            }
+        }
+
+        public ConnectionService ConnectionService { get; set; }
+
+        private Dictionary<string, ManualResetEvent> databaseAccessEvents = new Dictionary<string, ManualResetEvent>();
+        private object databaseAccessLock = new object();
+        public const int DefaultWaitToGetFullAccess = 60000;
+        public int waitToGetFullAccess = DefaultWaitToGetFullAccess;
+
+        private ManualResetEvent GetResetEvent(string serverName, string databaseName)
+        {
+            string key = GenerateKey(serverName, databaseName);
+            ManualResetEvent resetEvent = null;
+            lock (databaseAccessLock)
+            {
+                if (!databaseAccessEvents.TryGetValue(key, out resetEvent))
+                {
+                    resetEvent = new ManualResetEvent(true);
+                    databaseAccessEvents.Add(key, resetEvent);
+                }
+            }
+
+            return resetEvent;
+        }
+
+        public bool GainFullAccessToDatabase(string serverName, string databaseName)
+        {
+            /*
+            ManualResetEvent resetEvent = GetResetEvent(serverName, databaseName);
+            if (resetEvent.WaitOne(this.waitToGetFullAccess))
+            {
+                resetEvent.Reset();
+
+                foreach (IConnectedBindingQueue item in ConnectionService.ConnectedQueues)
+                {
+                    item.CloseConnections(serverName, databaseName);
+                }
+                return true;
+            }
+            else
+            {
+                throw new DatabaseFullAccessException($"Waited more than {waitToGetFullAccess} milli seconds for others to release the lock");
+            }
+            */
+            foreach (IConnectedBindingQueue item in ConnectionService.ConnectedQueues)
+            {
+                item.CloseConnections(serverName, databaseName);
+            }
+            return true;
+
+        }
+
+        public bool ReleaseAccess(string serverName, string databaseName)
+        {
+            /*
+            ManualResetEvent resetEvent = GetResetEvent(serverName, databaseName);
+
+            foreach (IConnectedBindingQueue item in ConnectionService.ConnectedQueues)
+            {
+                item.OpenConnections(serverName, databaseName);
+            }
+            
+            resetEvent.Set();
+            */
+            foreach (IConnectedBindingQueue item in ConnectionService.ConnectedQueues)
+            {
+                item.OpenConnections(serverName, databaseName);
+            }
+            return true;
+            
+        }
+
+        private string GenerateKey(string serverName, string databaseName)
+        {
+            return $"{serverName.ToLowerInvariant()}-{databaseName.ToLowerInvariant()}";
+        }
+
+        public void Dispose()
+        {
+            foreach (var resetEvent in databaseAccessEvents)
+            {
+                resetEvent.Value.Dispose();
+            }
+        }
+    }
+}

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

@@ -0,0 +1,40 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+namespace Microsoft.SqlTools.ServiceLayer.Connection
+{
+    /// <summary>
+    /// Any operation that needs full access to databas should implement this interface.
+    /// Make sure to call GainAccessToDatabase before the operation and ReleaseAccessToDatabase after
+    /// </summary>
+    public interface IFeatureWithFullDbAccess
+    {
+        /// <summary>
+        /// Database Lock Manager
+        /// </summary>
+        DatabaseLocksManager LockedDatabaseManager { get; set; }
+
+        /// <summary>
+        /// Makes sure the feature has fill access to the database
+        /// </summary>
+        bool GainAccessToDatabase();
+
+        /// <summary>
+        /// Release the access to db
+        /// </summary>
+        bool ReleaseAccessToDatabase();
+
+        /// <summary>
+        /// Server name
+        /// </summary>
+        string ServerName { get; }
+
+        /// <summary>
+        /// Database name
+        /// </summary>
+        string DatabaseName { get; }
+    }
+
+}