added contract to cancel restore plan and close the related connections (#522)

src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/Contracts/RestorePlanRequest.cs

@@ -87,4 +87,11 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts
             RequestType<RestoreParams, RestorePlanResponse> Type =
                 RequestType<RestoreParams, RestorePlanResponse>.Create("disasterrecovery/restoreplan");
     }
+
+    public class CancelRestorePlanRequest
+    {
+        public static readonly
+            RequestType<RestoreParams, bool> Type =
+                RequestType<RestoreParams, bool>.Create("disasterrecovery/cancelrestoreplan");
+    }
 }

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

@@ -14,6 +14,7 @@ using Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts;
 using Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation;
 using Microsoft.SqlTools.ServiceLayer.FileBrowser;
 using Microsoft.SqlTools.ServiceLayer.TaskServices;
+using Microsoft.SqlTools.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
 {
@@ -114,6 +115,9 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             // Create restore plan
             serviceHost.SetRequestHandler(RestorePlanRequest.Type, HandleRestorePlanRequest);
 
+            // Cancel restore plan
+            serviceHost.SetRequestHandler(CancelRestorePlanRequest.Type, HandleCancelRestorePlanRequest);
+
             // Create restore config
             serviceHost.SetRequestHandler(RestoreConfigInfoRequest.Type, HandleRestoreConfigInfoRequest);
 
@@ -164,6 +168,26 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             }
         }
 
+        /// <summary>
+        /// Handles a restore request
+        /// </summary>
+        internal async Task HandleCancelRestorePlanRequest(
+            RestoreParams restoreParams,
+            RequestContext<bool> requestContext)
+        {
+            bool result = false;
+            try
+            {
+                result = this.restoreDatabaseService.CancelRestorePlan(restoreParams);
+                await requestContext.SendResult(result);
+            }
+            catch (Exception ex)
+            {
+                Logger.Write(LogLevel.Error, "Failed to cancel restore session. error: " + ex.Message);
+                await requestContext.SendResult(result);
+            }
+        }
+
         /// <summary>
         /// Handles a restore request
         /// </summary>

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

@@ -66,6 +66,26 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             
         }
 
+        /// <summary>
+        /// Cancels existing restore plan
+        /// </summary>
+        public bool CancelRestorePlan(RestoreParams restoreParams)
+        {
+            RestoreDatabaseTaskDataObject restoreTaskObject = null;
+            string sessionId = restoreParams.SessionId;
+            if (!string.IsNullOrEmpty(sessionId) && sessions.TryGetValue(sessionId, out restoreTaskObject))
+            {
+                ServerConnection connection = restoreTaskObject?.Server?.ConnectionContext;
+                if (connection != null && connection.IsOpen)
+                {
+                    connection.Disconnect();
+                }
+                sessions.TryRemove(sessionId, out restoreTaskObject);
+                return true;
+            }
+            return false;
+        }
+
         /// <summary>
         /// Creates a restore plan, The result includes the information about the backup set, 
         /// the files and the database to restore to

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs

@@ -439,6 +439,40 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
             }
         }
 
+        [Fact]
+        public async Task CancelRestorePlanRequestShouldCancelSuccessfully()
+        {
+            await VerifyBackupFileCreated();
+
+            using (SelfCleaningTempFile queryTempFile = new SelfCleaningTempFile())
+            {
+                TestConnectionResult connectionResult = await LiveConnectionHelper.InitLiveConnectionInfoAsync("master", queryTempFile.FilePath);
+
+                string filePath = GetBackupFilePath(fullBackupFilePath);
+
+                RestoreParams restoreParams = new RestoreParams
+                {
+                    BackupFilePaths = filePath,
+                    OwnerUri = queryTempFile.FilePath
+                };
+
+                await RunAndVerify<RestorePlanResponse>(
+                    test: (requestContext) => service.HandleRestorePlanRequest(restoreParams, requestContext),
+                    verify: ((result) =>
+                    {
+                        restoreParams.SessionId = result.SessionId;
+                        Assert.True(result.DbFiles.Any());
+                    }));
+
+                await RunAndVerify<bool>(
+                   test: (requestContext) => service.HandleCancelRestorePlanRequest(restoreParams, requestContext),
+                   verify: ((result) =>
+                   {
+                       Assert.True(result);
+                   }));
+            }
+        }
+
         [Fact]
         public async Task RestoreConfigInfoRequestShouldReturnResponse()
         {