diff --git a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/Contracts/RestorePlanRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/Contracts/RestorePlanRequest.cs index b7987407..31623796 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/Contracts/RestorePlanRequest.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/Contracts/RestorePlanRequest.cs @@ -87,4 +87,11 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts RequestType Type = RequestType.Create("disasterrecovery/restoreplan"); } + + public class CancelRestorePlanRequest + { + public static readonly + RequestType Type = + RequestType.Create("disasterrecovery/cancelrestoreplan"); + } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs index 8557c738..a7eb6f12 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.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 } } + /// + /// Handles a restore request + /// + internal async Task HandleCancelRestorePlanRequest( + RestoreParams restoreParams, + RequestContext 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); + } + } + /// /// Handles a restore request /// diff --git a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs index 17a7f6ac..d2f2d0d7 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs @@ -66,6 +66,26 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation } + /// + /// Cancels existing restore plan + /// + 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; + } + /// /// Creates a restore plan, The result includes the information about the backup set, /// the files and the database to restore to diff --git a/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs b/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs index 84153adc..9bfee231 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs +++ b/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( + test: (requestContext) => service.HandleRestorePlanRequest(restoreParams, requestContext), + verify: ((result) => + { + restoreParams.SessionId = result.SessionId; + Assert.True(result.DbFiles.Any()); + })); + + await RunAndVerify( + test: (requestContext) => service.HandleCancelRestorePlanRequest(restoreParams, requestContext), + verify: ((result) => + { + Assert.True(result); + })); + } + } + [Fact] public async Task RestoreConfigInfoRequestShouldReturnResponse() {