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/Admin/AdminService.cs

@@ -7,7 +7,6 @@ using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Admin.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Hosting;
-using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using System;
 using System.Threading.Tasks;
 using System.Xml;
@@ -28,8 +27,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Admin
         private static readonly ConcurrentDictionary<string, DatabaseTaskHelper> serverTaskHelperMap =
             new ConcurrentDictionary<string, DatabaseTaskHelper>();
 
-        private static DatabaseTaskHelper taskHelper;
-
         /// <summary>
         /// Default, parameterless constructor.
         /// </summary>
@@ -91,13 +88,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Admin
                     optionsParams.OwnerUri,
                     out connInfo);
 
-                if (taskHelper == null)
+                using (var taskHelper = CreateDatabaseTaskHelper(connInfo))
                 {
-                    taskHelper = CreateDatabaseTaskHelper(connInfo);
+                    response.DefaultDatabaseInfo = DatabaseTaskHelper.DatabasePrototypeToDatabaseInfo(taskHelper.Prototype);
+                    await requestContext.SendResult(response);
                 }
-
-                response.DefaultDatabaseInfo = DatabaseTaskHelper.DatabasePrototypeToDatabaseInfo(taskHelper.Prototype);
-                await requestContext.SendResult(response);
             }
             catch (Exception ex)
             {
@@ -120,25 +115,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Admin
                     databaseParams.OwnerUri,
                     out connInfo);
 
-                if (taskHelper == null)
+                using (var taskHelper = CreateDatabaseTaskHelper(connInfo))
                 {
-                    taskHelper = CreateDatabaseTaskHelper(connInfo);
+                    DatabasePrototype prototype = taskHelper.Prototype;
+                    DatabaseTaskHelper.ApplyToPrototype(databaseParams.DatabaseInfo, taskHelper.Prototype);
+
+                    Database db = prototype.ApplyChanges();
+
+                    await requestContext.SendResult(new CreateDatabaseResponse()
+                    {
+                        Result = true,
+                        TaskId = 0
+                    });
                 }
-
-                DatabasePrototype prototype = taskHelper.Prototype;
-                DatabaseTaskHelper.ApplyToPrototype(databaseParams.DatabaseInfo, taskHelper.Prototype);
-
-                Database db = prototype.ApplyChanges();
-                if (db != null)
-                {
-                    taskHelper = null;
-                }
-
-                await requestContext.SendResult(new CreateDatabaseResponse()
-                {
-                    Result = true,
-                    TaskId = 0
-                }); 
             }
             catch (Exception ex)
             {
@@ -182,9 +171,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Admin
         /// <param name="connInfo"></param>
         /// <returns></returns>
         internal static DatabaseInfo GetDatabaseInfo(ConnectionInfo connInfo)
-        {            
-            DatabaseTaskHelper taskHelper = CreateDatabaseTaskHelper(connInfo, true);
-            return DatabaseTaskHelper.DatabasePrototypeToDatabaseInfo(taskHelper.Prototype);
+        {
+            using (DatabaseTaskHelper taskHelper = CreateDatabaseTaskHelper(connInfo, true))
+            {
+                return DatabaseTaskHelper.DatabasePrototypeToDatabaseInfo(taskHelper.Prototype);
+            }
         }
 
         /// <summary>
@@ -205,6 +196,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Admin
                 : string.Format("{0},{1}", connectionDetails.ServerName, connectionDetails.Port.Value);
 
             // check if the connection is using SQL Auth or Integrated Auth
+            //TODO: ConnectionQueue try to get an existing connection (ConnectionQueue)
             if (string.Equals(connectionDetails.AuthenticationType, "SqlLogin", StringComparison.OrdinalIgnoreCase))
             {
                 var passwordSecureString = BuildSecureStringFromPassword(connectionDetails.Password);

src/Microsoft.SqlTools.ServiceLayer/Admin/Database/DatabaseTaskHelper.cs

@@ -6,6 +6,7 @@
 using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlServer.Management.Smo;
 using Microsoft.SqlTools.ServiceLayer.Admin.Contracts;
+using Microsoft.SqlTools.Utility;
 using System;
 using System.Collections;
 using System.Collections.Generic;
@@ -15,7 +16,7 @@ using System.Xml;
 
 namespace Microsoft.SqlTools.ServiceLayer.Admin
 {
-    public class DatabaseTaskHelper
+    public class DatabaseTaskHelper: IDisposable
     {
         private DatabasePrototype prototype;
 
@@ -184,5 +185,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Admin
             }
             return prototype;
         }
+
+        public void Dispose()
+        {
+            try
+            {
+                if (this.DataContainer != null)
+                {
+                    this.DataContainer.Dispose();
+                }
+            }
+            catch(Exception ex)
+            {
+                Logger.Write(LogLevel.Warning, $"Failed to disconnect Database task Helper connection. Error: {ex.Message}");
+            }
+        }
     }
 }

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

src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/BackupOperation/BackupOperation.cs

@@ -313,6 +313,17 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             {
                 throw;
             }
+            finally
+            {
+                if (this.serverConnection != null)
+                {
+                    this.serverConnection.Disconnect();
+                }
+                if(this.dataContainer != null)
+                {
+                    this.dataContainer.Dispose();
+                }
+            }
         }
 
         /// <summary>

src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs

@@ -142,13 +142,17 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
 
                 if (connInfo != null)
                 {
-                    DatabaseTaskHelper helper = AdminService.CreateDatabaseTaskHelper(connInfo, databaseExists: true);
-                    SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connInfo);
-                    if (sqlConn != null && !connInfo.IsSqlDW && !connInfo.IsAzure)
+                    using (DatabaseTaskHelper helper = AdminService.CreateDatabaseTaskHelper(connInfo, databaseExists: true))
                     {
-                        BackupConfigInfo backupConfigInfo = this.GetBackupConfigInfo(helper.DataContainer, sqlConn, sqlConn.Database);
-                        backupConfigInfo.DatabaseInfo = AdminService.GetDatabaseInfo(connInfo);
-                        response.BackupConfigInfo = backupConfigInfo;
+                        using (SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connInfo))
+                        {
+                            if (sqlConn != null && !connInfo.IsSqlDW && !connInfo.IsAzure)
+                            {
+                                BackupConfigInfo backupConfigInfo = this.GetBackupConfigInfo(helper.DataContainer, sqlConn, sqlConn.Database);
+                                backupConfigInfo.DatabaseInfo = AdminService.GetDatabaseInfo(connInfo);
+                                response.BackupConfigInfo = backupConfigInfo;
+                            }
+                        }
                     }
                 }
 
@@ -233,7 +237,6 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             RequestContext<RestoreResponse> requestContext)
         {
             RestoreResponse response = new RestoreResponse();
-
             try
             {
                 ConnectionInfo connInfo;
@@ -243,10 +246,12 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                 {
                     try
                     {
-                        RestoreDatabaseTaskDataObject restoreDataObject = this.restoreDatabaseService.CreateRestoreDatabaseTaskDataObject(restoreParams);
+                       
+                        RestoreDatabaseTaskDataObject restoreDataObject = this.restoreDatabaseService.CreateRestoreDatabaseTaskDataObject(restoreParams, connInfo);
 
                         if (restoreDataObject != null)
                         {
+                            restoreDataObject.LockedDatabaseManager = ConnectionServiceInstance.LockedDatabaseManager;
                             // create task metadata
                             TaskMetadata metadata = TaskMetadata.Create(restoreParams, SR.RestoreTaskName, restoreDataObject, ConnectionServiceInstance);
                             metadata.DatabaseName = restoreParams.TargetDatabaseName;
@@ -297,6 +302,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                 {
                     DatabaseTaskHelper helper = AdminService.CreateDatabaseTaskHelper(connInfo, databaseExists: true);
                     SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connInfo);
+                    // Connection gets discounnected when backup is done
 
                     BackupOperation backupOperation = CreateBackupOperation(helper.DataContainer, sqlConn, backupParams.BackupInfo);
                     SqlTask sqlTask = null;
@@ -332,17 +338,19 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
 
                 if (connInfo != null)
                 {
-                    sqlConn = ConnectionService.OpenSqlConnection(connInfo);
-                    if (sqlConn != null && !connInfo.IsSqlDW && !connInfo.IsAzure)
+                    using (sqlConn = ConnectionService.OpenSqlConnection(connInfo))
                     {
-                        connectionInfo = connInfo;
-                        return true;
+                        if (sqlConn != null && !connInfo.IsSqlDW && !connInfo.IsAzure)
+                        {
+                            connectionInfo = connInfo;
+                            return true;
+                        }
                     }
                 }
             }
             catch
             {
-                if(sqlConn != null)
+                if(sqlConn != null && sqlConn.State == System.Data.ConnectionState.Open)
                 {
                     sqlConn.Close();
                 }

src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs

@@ -88,10 +88,6 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                             RestoreAsFileName = x.PhysicalNameRelocate
                         });
                         response.CanRestore = CanRestore(restoreDataObject);
-                        if (!response.CanRestore)
-                        {
-                            response.ErrorMessage = SR.NoBackupsetsToRestore;
-                        }
 
                         response.PlanDetails.Add(LastBackupTaken, 
                             RestorePlanDetailInfo.Create(name: LastBackupTaken, currentValue: restoreDataObject.GetLastBackupTaken(), isReadOnly: true));
@@ -150,7 +146,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// </summary>
         /// <param name="restoreParams">Restore request parameters</param>
         /// <returns>Restore task object</returns>
-        public RestoreDatabaseTaskDataObject CreateRestoreDatabaseTaskDataObject(RestoreParams restoreParams)
+        public RestoreDatabaseTaskDataObject CreateRestoreDatabaseTaskDataObject(RestoreParams restoreParams, ConnectionInfo connectionInfo = null)
         {
             RestoreDatabaseTaskDataObject restoreTaskObject = null;
             string sessionId = string.IsNullOrWhiteSpace(restoreParams.SessionId) ? Guid.NewGuid().ToString() : restoreParams.SessionId;
@@ -161,6 +157,10 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             }
             restoreTaskObject.SessionId = sessionId;
             restoreTaskObject.RestoreParams = restoreParams;
+            if (connectionInfo != null)
+            {
+                restoreTaskObject.ConnectionInfo = connectionInfo;
+            }
             
             return restoreTaskObject;
         }

src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseTaskDataObject.cs

@@ -14,6 +14,7 @@ using Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts;
 using Microsoft.SqlTools.ServiceLayer.TaskServices;
 using Microsoft.SqlTools.ServiceLayer.Utility;
 using Microsoft.SqlTools.Utility;
+using Microsoft.SqlTools.ServiceLayer.Connection;
 
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 {
@@ -65,7 +66,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
     /// <summary>
     /// Includes the plan with all the data required to do a restore operation on server
     /// </summary>
-    public class RestoreDatabaseTaskDataObject : SmoScriptableTaskOperation, IRestoreDatabaseTaskDataObject
+    public class RestoreDatabaseTaskDataObject : SmoScriptableOperationWithFullDbAccess, IRestoreDatabaseTaskDataObject
     {
 
         private const char BackupMediaNameSeparator = ',';
@@ -266,29 +267,66 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             base.Execute(mode);
         }
 
+        public ConnectionInfo ConnectionInfo { get; set; }
+
+        public override string ServerName
+        {
+            get
+            {
+                if (this.ConnectionInfo != null)
+                {
+                    return this.ConnectionInfo.ConnectionDetails.ServerName;
+                }
+
+                return this.Server.Name;
+            }
+        }
+
+        public override string DatabaseName
+        {
+            get
+            {
+                return TargetDatabaseName;
+            }
+        }
+
         /// <summary>
         /// Executes the restore operations
         /// </summary>
         public override void Execute()
         {
-            if (IsValid && RestorePlan.RestoreOperations != null && RestorePlan.RestoreOperations.Any())
+            try
             {
-                // Restore Plan should be already created and updated at this point
-
-                RestorePlan restorePlan = GetRestorePlanForExecutionAndScript();
-
-                if (restorePlan != null && restorePlan.RestoreOperations.Count > 0)
+                if (IsValid && RestorePlan.RestoreOperations != null && RestorePlan.RestoreOperations.Any())
                 {
-                    restorePlan.PercentComplete += (object sender, PercentCompleteEventArgs e) =>
+                    // Restore Plan should be already created and updated at this point
+
+                    RestorePlan restorePlan = GetRestorePlanForExecutionAndScript();
+
+                    if (restorePlan != null && restorePlan.RestoreOperations.Count > 0)
                     {
-                        OnMessageAdded(new TaskMessage { Description = $"{e.Percent}%", Status = SqlTaskStatus.InProgress });
-                    };
-                    restorePlan.Execute();
+                        restorePlan.PercentComplete += (object sender, PercentCompleteEventArgs e) =>
+                        {
+                            OnMessageAdded(new TaskMessage { Description = $"{e.Percent}%", Status = SqlTaskStatus.InProgress });
+                        };
+                        restorePlan.Execute();
+                    }
+                }
+                else
+                {
+                    throw new InvalidOperationException(SR.RestoreNotSupported);
                 }
             }
-            else
+            catch(Exception ex)
             {
-                throw new InvalidOperationException(SR.RestoreNotSupported);
+                throw ex;
+            }
+            finally
+            {
+                if (this.Server.ConnectionContext.IsOpen)
+                {
+                    this.Server.ConnectionContext.Disconnect();
+                }
             }
         }
 

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/BindingQueue.cs

@@ -8,6 +8,7 @@ using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.Utility;
+using System.Linq;
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 {    
@@ -112,6 +113,20 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             }      
         }
 
+        protected IEnumerable<IBindingContext> GetBindingContexts(string keyPrefix)
+        {
+            // use a default binding context for disconnected requests
+            if (string.IsNullOrWhiteSpace(keyPrefix))
+            {
+                keyPrefix = "disconnected_binding_context";
+            }
+
+            lock (this.bindingContextLock)
+            {
+                return this.BindingContextMap.Where(x => x.Key.StartsWith(keyPrefix)).Select(v => v.Value);
+            }
+        }
+
         /// <summary>
         /// Checks if a binding context already exists for the provided context key
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/ConnectedBindingQueue.cs

@@ -13,13 +13,28 @@ using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
+using System.Threading;
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 {
+    public interface IConnectedBindingQueue
+    {
+        void CloseConnections(string serverName, string databaseName);
+        void OpenConnections(string serverName, string databaseName);
+        string AddConnectionContext(ConnectionInfo connInfo, bool overwrite = false);
+        void Dispose();
+        QueueItem QueueBindingOperation(
+            string key,
+            Func<IBindingContext, CancellationToken, object> bindOperation,
+            Func<IBindingContext, object> timeoutOperation = null,
+            int? bindingTimeout = null,
+            int? waitForLockTimeout = null);
+    }
+
     /// <summary>
     /// ConnectedBindingQueue class for processing online binding requests
     /// </summary>
-    public class ConnectedBindingQueue : BindingQueue<ConnectedBindingContext>
+    public class ConnectedBindingQueue : BindingQueue<ConnectedBindingContext>, IConnectedBindingQueue
     {
         internal const int DefaultBindingTimeout = 500;
 
@@ -64,6 +79,44 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             );
         }
 
+        /// <summary>
+        /// Generate a unique key based on the ConnectionInfo object
+        /// </summary>
+        /// <param name="connInfo"></param>
+        private string GetConnectionContextKey(string serverName, string databaseName)
+        {
+            return string.Format("{0}_{1}",
+                serverName ?? "NULL",
+                databaseName ?? "NULL");
+            
+        }
+
+        public void CloseConnections(string serverName, string databaseName)
+        {
+            string connectionKey = GetConnectionContextKey(serverName, databaseName);
+            var contexts = GetBindingContexts(connectionKey);
+            foreach (var bindingContext in contexts)
+            {
+                if (bindingContext.BindingLock.WaitOne(2000))
+                {
+                    bindingContext.ServerConnection.Disconnect();
+                }
+            }
+        }
+
+        public void OpenConnections(string serverName, string databaseName)
+        {
+            string connectionKey = GetConnectionContextKey(serverName, databaseName);
+            var contexts = GetBindingContexts(connectionKey);
+            foreach (var bindingContext in contexts)
+            {
+                if (bindingContext.BindingLock.WaitOne(2000))
+                {
+                    //bindingContext.ServerConnection.Connect();
+                }
+            }
+        }
+
         /// <summary>
         /// Use a ConnectionInfo item to create a connected binding context
         /// </summary>
@@ -98,7 +151,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 {
                     bindingContext.BindingLock.Reset();
                     SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connInfo);
-
+                   
                     // populate the binding context to work with the SMO metadata provider
                     bindingContext.ServerConnection = new ServerConnection(sqlConn);
 

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs

@@ -150,6 +150,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 if (connectionService == null)
                 {
                     connectionService = ConnectionService.Instance;
+                    connectionService.RegisterConnectedQueue("LanguageService", bindingQueue);
                 }
                 return connectionService;
             }

src/Microsoft.SqlTools.ServiceLayer/Metadata/MetadataService.cs

@@ -129,11 +129,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Metadata
                 ColumnMetadata[] metadata = null;
                 if (connInfo != null) 
                 {
-                    SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connInfo);                    
-                    TableMetadata table = new SmoMetadataFactory().GetObjectMetadata(
-                        sqlConn, metadataParams.Schema, 
-                        metadataParams.ObjectName, objectType);
-                    metadata = table.Columns;               
+                    using (SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connInfo))
+                    {
+                        TableMetadata table = new SmoMetadataFactory().GetObjectMetadata(
+                            sqlConn, metadataParams.Schema,
+                            metadataParams.ObjectName, objectType);
+                        metadata = table.Columns;
+                    }
                 }
 
                 await requestContext.SendResult(new TableMetadataResult

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs

@@ -328,7 +328,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
                                 {
                                     children.Add(item);
                                     item.Parent = this;
-
                                 }
                             }
                         }

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs

@@ -25,6 +25,7 @@ using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Utility;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlTools.Utility;
+using Microsoft.SqlServer.Management.Common;
 
 namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
 {
@@ -61,6 +62,18 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             applicableNodeChildFactories = new Lazy<Dictionary<string, HashSet<ChildFactory>>>(() => PopulateFactories());
         }
 
+        internal ConnectedBindingQueue ConnectedBindingQueue
+        {
+            get
+            {
+                return bindingQueue;
+            }
+            set
+            {
+                this.bindingQueue = value;
+            }
+        }
+
         /// <summary>
         /// Internal for testing only
         /// </summary>
@@ -99,6 +112,15 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             Validate.IsNotNull(nameof(provider), provider);
             serviceProvider = provider;
             connectionService = provider.GetService<ConnectionService>();
+            try
+            {
+                connectionService.RegisterConnectedQueue("OE", bindingQueue);
+
+            }
+            catch(Exception ex)
+            {
+                Logger.Write(LogLevel.Error, ex.Message);
+            }
         }
 
         /// <summary>
@@ -119,6 +141,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             {
                 workspaceService.RegisterConfigChangeCallback(HandleDidChangeConfigurationNotification);
             }
+
         }
 
         /// <summary>
@@ -369,7 +392,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                 try
                 {
                     QueueItem queueItem = bindingQueue.QueueBindingOperation(
-                           key: session.Uri,
+                           key: bindingQueue.AddConnectionContext(session.ConnectionInfo),
                            bindingTimeout: PrepopulateBindTimeout,
                            waitForLockTimeout: PrepopulateBindTimeout,
                            bindOperation: (bindingContext, cancelToken) =>
@@ -413,19 +436,40 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         {
             try
             {
-                ObjectExplorerSession session;
+                ObjectExplorerSession session = null;
                 connectionDetails.PersistSecurityInfo = true;
-                ConnectParams connectParams = new ConnectParams() { OwnerUri = uri, Connection = connectionDetails };
+                ConnectParams connectParams = new ConnectParams() { OwnerUri = uri, Connection = connectionDetails, Type = Connection.ConnectionType.ObjectExplorer };
 
+                ConnectionInfo connectionInfo;
                 ConnectionCompleteParams connectionResult = await Connect(connectParams, uri);
+                if (!connectionService.TryFindConnection(uri, out connectionInfo))
+                {
+                    return null;
+                }
+
                 if (connectionResult == null)
                 {
                     // Connection failed and notification is already sent
                     return null;
                 }
 
-                session = ObjectExplorerSession.CreateSession(connectionResult, serviceProvider);
-                sessionMap.AddOrUpdate(uri, session, (key, oldSession) => session);
+                QueueItem queueItem = bindingQueue.QueueBindingOperation(
+                           key: bindingQueue.AddConnectionContext(connectionInfo),
+                           bindingTimeout: PrepopulateBindTimeout,
+                           waitForLockTimeout: PrepopulateBindTimeout,
+                           bindOperation: (bindingContext, cancelToken) =>
+                           {
+                               session = ObjectExplorerSession.CreateSession(connectionResult, serviceProvider, bindingContext.ServerConnection);
+                               session.ConnectionInfo = connectionInfo;
+
+                               sessionMap.AddOrUpdate(uri, session, (key, oldSession) => session);
+                               return session;
+                           });
+                queueItem.ItemProcessed.WaitOne();
+                if (queueItem.GetResultAsT<ObjectExplorerSession>() != null)
+                {
+                    session = queueItem.GetResultAsT<ObjectExplorerSession>();
+                }
                 return session;
             }
             catch(Exception ex)
@@ -657,11 +701,13 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             public string Uri { get; private set; }
             public TreeNode Root { get; private set; }
 
+            public ConnectionInfo ConnectionInfo { get; set; }
+
             public string ErrorMessage { get; set; }
 
-            public static ObjectExplorerSession CreateSession(ConnectionCompleteParams response, IMultiServiceProvider serviceProvider)
+            public static ObjectExplorerSession CreateSession(ConnectionCompleteParams response, IMultiServiceProvider serviceProvider, ServerConnection serverConnection)
             {
-                ServerNode rootNode = new ServerNode(response, serviceProvider);
+                ServerNode rootNode = new ServerNode(response, serviceProvider, serverConnection);
                 var session = new ObjectExplorerSession(response.OwnerUri, rootNode, serviceProvider, serviceProvider.GetService<ConnectionService>());
                 if (!DatabaseUtils.IsSystemDatabaseConnection(response.ConnectionSummary.DatabaseName))
                 {

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/ServerNode.cs

@@ -27,13 +27,12 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
     {
         private ConnectionSummary connectionSummary;
         private ServerInfo serverInfo;
-        private string connectionUri;
         private Lazy<SmoQueryContext> context;
-        private ConnectionService connectionService;
         private SmoWrapper smoWrapper;
         private SqlServerType sqlServerType;
+        private ServerConnection serverConnection;
 
-        public ServerNode(ConnectionCompleteParams connInfo, IMultiServiceProvider serviceProvider)
+        public ServerNode(ConnectionCompleteParams connInfo, IMultiServiceProvider serviceProvider, ServerConnection serverConnection)
             : base()
         {
             Validate.IsNotNull(nameof(connInfo), connInfo);
@@ -42,12 +41,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
 
             this.connectionSummary = connInfo.ConnectionSummary;
             this.serverInfo = connInfo.ServerInfo;
-            this.connectionUri = connInfo.OwnerUri;
             this.sqlServerType = ServerVersionHelper.CalculateServerType(this.serverInfo);
 
-            this.connectionService = serviceProvider.GetService<ConnectionService>();
-
             this.context = new Lazy<SmoQueryContext>(() => CreateContext(serviceProvider));
+            this.serverConnection = serverConnection;
 
             NodeValue = connectionSummary.ServerName;
             IsAlwaysLeaf = false;
@@ -130,43 +127,22 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
         private SmoQueryContext CreateContext(IMultiServiceProvider serviceProvider)
         {
             string exceptionMessage;
-            ConnectionInfo connectionInfo;
-            SqlConnection connection = null;
-            // Get server object from connection
-            if (!connectionService.TryFindConnection(this.connectionUri, out connectionInfo) ||
-                connectionInfo.AllConnections == null || connectionInfo.AllConnections.Count == 0)
-            {
-                ErrorStateMessage = string.Format(CultureInfo.CurrentCulture, 
-                    SR.ServerNodeConnectionError, connectionSummary.ServerName);
-                return null;
-            }
-            //TODO: figure out how to use existing connections
-            DbConnection dbConnection = connectionInfo.AllConnections.First();
-            ReliableSqlConnection reliableSqlConnection = dbConnection as ReliableSqlConnection;
-            SqlConnection sqlConnection = dbConnection as SqlConnection;
-            if (reliableSqlConnection != null)
-            {
-                connection = reliableSqlConnection.GetUnderlyingConnection();
-            }
-            else if (sqlConnection != null)
-            {
-                connection = sqlConnection;
-            }
-            else
-            {
-                ErrorStateMessage = string.Format(CultureInfo.CurrentCulture,
-                   SR.ServerNodeConnectionError, connectionSummary.ServerName);
-                return null;
-            }
-
+   
             try
             {
-                Server server = SmoWrapper.CreateServer(connection);
-                return new SmoQueryContext(server, serviceProvider, SmoWrapper)
+                Server server = SmoWrapper.CreateServer(this.serverConnection);
+                if (server != null)
                 {
-                    Parent = server,
-                    SqlServerType = this.sqlServerType
-                };
+                    return new SmoQueryContext(server, serviceProvider, SmoWrapper)
+                    {
+                        Parent = server,
+                        SqlServerType = this.sqlServerType
+                    };
+                }
+                else
+                {
+                    return null;
+                }
             }
             catch (ConnectionFailureException cfe)
             {

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoQueryContext.cs

@@ -174,7 +174,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
         /// the only way to easily access is via the server object. This should be called during access of
         /// any of the object properties
         /// </summary>
-        private void EnsureConnectionOpen(SmoObjectBase smoObj)
+        public void EnsureConnectionOpen(SmoObjectBase smoObj)
         {
             if (!smoWrapper.IsConnectionOpen(smoObj))
             {

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoWrapper.cs

@@ -15,10 +15,9 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
     /// </summary>
     internal class SmoWrapper
     {
-        public virtual Server CreateServer(SqlConnection connection)
+        public virtual Server CreateServer(ServerConnection serverConn)
         {
-            ServerConnection serverConn = new ServerConnection(connection);
-            return new Server(serverConn);
+            return serverConn == null ? null : new Server(serverConn);
         }
 
         public virtual bool IsConnectionOpen(SmoObjectBase smoObj)

src/Microsoft.SqlTools.ServiceLayer/Profiler/ProfilerSessionMonitor.cs

@@ -15,6 +15,7 @@ using Microsoft.SqlServer.Management.Sdk.Sfc;
 using Microsoft.SqlServer.Management.XEvent;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Profiler.Contracts;
+using Microsoft.SqlTools.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.Profiler
 {
@@ -150,6 +151,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
                     }
                 }
             }
+            catch(Exception ex)
+            {
+                Logger.Write(LogLevel.Warning, "Failed to pool session. error: " + ex.Message);
+            }
             finally
             {
                 session.IsPolling = false;

src/Microsoft.SqlTools.ServiceLayer/TaskServices/SmoScriptableOperationWithFullDbAccess.cs

@@ -0,0 +1,102 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlServer.Management.Smo;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.Utility;
+
+namespace Microsoft.SqlTools.ServiceLayer.TaskServices
+{
+    public abstract class SmoScriptableOperationWithFullDbAccess : SmoScriptableTaskOperation, IFeatureWithFullDbAccess
+    {
+        private DatabaseLocksManager lockedDatabaseManager;
+        /// <summary>
+        /// If an error occurred during task execution, this field contains the error message text
+        /// </summary>
+        public override abstract string ErrorMessage { get; }
+
+        /// <summary>
+        /// SMO Server instance used for the operation
+        /// </summary>
+        public override abstract Server Server { get; }
+        public DatabaseLocksManager LockedDatabaseManager
+        {
+            get
+            {
+                if (lockedDatabaseManager == null)
+                {
+                    lockedDatabaseManager = ConnectionService.Instance.LockedDatabaseManager;
+                }
+                return lockedDatabaseManager;
+            }
+            set
+            {
+                lockedDatabaseManager = value;
+            }
+        }
+
+        public abstract string ServerName { get; }
+
+        public abstract string DatabaseName { get; }
+
+        /// <summary>
+        /// Cancels the operation
+        /// </summary>
+        public override abstract void Cancel();
+
+        /// <summary>
+        /// Executes the operations
+        /// </summary>
+        public override abstract void Execute();
+
+        /// <summary>
+        /// Execute the operation for given execution mode
+        /// </summary>
+        /// <param name="mode"></param>
+        public override void Execute(TaskExecutionMode mode)
+        {
+            bool hasAccessToDb = false;
+            try
+            {
+                hasAccessToDb = GainAccessToDatabase();
+                base.Execute(mode);
+            }
+            catch (DatabaseFullAccessException databaseFullAccessException)
+            {
+                Logger.Write(LogLevel.Warning, $"Failed to gain access to database. server|database:{ServerName}|{DatabaseName}");
+                throw databaseFullAccessException;
+            }
+            catch
+            {
+                throw;
+            }
+            finally
+            {
+                if (hasAccessToDb)
+                {
+                    ReleaseAccessToDatabase();
+                }
+            }
+        }
+
+        public bool GainAccessToDatabase()
+        {
+            if (LockedDatabaseManager != null)
+            {
+                return LockedDatabaseManager.GainFullAccessToDatabase(ServerName, DatabaseName);
+            }
+            return false;
+        }
+
+        public bool ReleaseAccessToDatabase()
+        {
+            if (LockedDatabaseManager != null)
+            {
+                return LockedDatabaseManager.ReleaseAccess(ServerName, DatabaseName);
+            }
+            return false;
+        }
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/TaskServices/SmoScriptableTaskOperation.cs

@@ -76,43 +76,15 @@ namespace Microsoft.SqlTools.ServiceLayer.TaskServices
             var currentExecutionMode = Server.ConnectionContext.SqlExecutionModes;
             try
             {
-                
                 if (Server != null)
                 {
                     Server.ConnectionContext.CapturedSql.Clear();
-                    switch (mode)
-                    {
-                        case TaskExecutionMode.Execute:
-                            {
-                                Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.ExecuteSql;
-                                break;
-                            }
-                        case TaskExecutionMode.ExecuteAndScript:
-                            {
-                                Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.ExecuteAndCaptureSql;
-                                break;
-                            }
-                        case TaskExecutionMode.Script:
-                            {
-                                Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.CaptureSql;
-                                break;
-                            }
-                    }
+                    SetExecutionMode(mode);
                 }
 
                 Execute();
-                if (mode == TaskExecutionMode.Script || mode == TaskExecutionMode.ExecuteAndScript)
-                {
-                    this.ScriptContent = GetScriptContent();
-                    if (SqlTask != null)
-                    {
-                        OnScriptAdded(new TaskScript
-                        {
-                            Status = SqlTaskStatus.Succeeded,
-                            Script = this.ScriptContent
-                        });
-                    }
-                }
+                GenerateScript(mode);
+                
             }
             catch
             {
@@ -126,6 +98,44 @@ namespace Microsoft.SqlTools.ServiceLayer.TaskServices
 
         }
 
+        protected void GenerateScript(TaskExecutionMode mode)
+        {
+            if (mode == TaskExecutionMode.Script || mode == TaskExecutionMode.ExecuteAndScript)
+            {
+                this.ScriptContent = GetScriptContent();
+                if (SqlTask != null)
+                {
+                    OnScriptAdded(new TaskScript
+                    {
+                        Status = SqlTaskStatus.Succeeded,
+                        Script = this.ScriptContent
+                    });
+                }
+            }
+        }
+
+        protected void SetExecutionMode(TaskExecutionMode mode)
+        {
+            switch (mode)
+            {
+                case TaskExecutionMode.Execute:
+                    {
+                        Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.ExecuteSql;
+                        break;
+                    }
+                case TaskExecutionMode.ExecuteAndScript:
+                    {
+                        Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.ExecuteAndCaptureSql;
+                        break;
+                    }
+                case TaskExecutionMode.Script:
+                    {
+                        Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.CaptureSql;
+                        break;
+                    }
+            }
+        }
+
         private string GetScriptContent()
         {
             StringBuilder sb = new StringBuilder();