From 44d3e4b17e109f847ffbba4547828eafcb94bd4e Mon Sep 17 00:00:00 2001
From: Kate Shin <hyoshi@microsoft.com>
Date: Thu, 31 Aug 2017 11:00:02 -0700
Subject: [PATCH] Close connection after backup execution (#449)

---
 .../BackupOperation/BackupOperation.cs             |  2 --
 .../DisasterRecovery/DisasterRecoveryService.cs    | 14 ++++++++++++++
 .../TaskServices/TaskMetadata.cs                   |  7 +++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/BackupOperation/BackupOperation.cs b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/BackupOperation/BackupOperation.cs
index 645d4be8..d5936338 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/BackupOperation/BackupOperation.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/BackupOperation/BackupOperation.cs
@@ -308,8 +308,6 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                         }
                     }
                 }
-
-                
             }
             catch(Exception)
             {
diff --git a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs
index 6c14f285..b33df78f 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs
@@ -14,6 +14,7 @@ using Microsoft.SqlTools.ServiceLayer.TaskServices;
 using System.Threading;
 using Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation;
 using System.Globalization;
+using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
 {
@@ -281,6 +282,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                     TaskMetadata metadata = TaskMetadata.Create(backupParams, SR.BackupTaskName, backupOperation, ConnectionServiceInstance);
                    
                     sqlTask = SqlTaskManagerInstance.CreateAndRun<SqlTask>(metadata);
+                    sqlTask.StatusChanged += CloseConnection;
                 }
                 else
                 {
@@ -352,6 +354,18 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             return false;
         }
 
+        private void CloseConnection(object sender, TaskEventArgs<SqlTaskStatus> e)
+        {
+            SqlTask sqlTask = e.SqlTask;
+            if (sqlTask != null && sqlTask.IsCompleted)
+            {
+                connectionService.Disconnect(new DisconnectParams()
+                {
+                    OwnerUri = sqlTask.TaskMetadata.OwnerUri
+                });
+            }
+        }
+
         private BackupOperation CreateBackupOperation(CDataContainer dataContainer, SqlConnection sqlConnection)
         {
             BackupOperation backupOperation = new BackupOperation();
diff --git a/src/Microsoft.SqlTools.ServiceLayer/TaskServices/TaskMetadata.cs b/src/Microsoft.SqlTools.ServiceLayer/TaskServices/TaskMetadata.cs
index 6ec143a5..c1143a69 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/TaskServices/TaskMetadata.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/TaskServices/TaskMetadata.cs
@@ -46,6 +46,11 @@ namespace Microsoft.SqlTools.ServiceLayer.TaskServices
         /// </summary>
         public ITaskOperation TaskOperation { get; set; }
 
+        /// <summary>
+        /// Connection uri
+        /// </summary>
+        public string OwnerUri { get; set; }
+
         /// <summary>
         /// Creates task metadata given the request parameters
         /// </summary>
@@ -84,6 +89,8 @@ namespace Microsoft.SqlTools.ServiceLayer.TaskServices
             taskMetadata.TaskExecutionMode = scriptableRequestParams.TaskExecutionMode;
 
             taskMetadata.TaskOperation = taskOperation;
+            taskMetadata.OwnerUri = requestParam.OwnerUri;
+
             return taskMetadata;
         }