From 09dd0579d09ac3d94369595f91da2331e67cdb90 Mon Sep 17 00:00:00 2001
From: Leila Lali <llali@microsoft.com>
Date: Wed, 25 Oct 2017 10:23:28 -0700
Subject: [PATCH] Fix/restore scripting issue (#521)

* fixed the bug with setting server execution mode before creating restore plan
---
 .../RestoreDatabaseTaskDataObject.cs              | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseTaskDataObject.cs b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseTaskDataObject.cs
index 391d1151..3e9ea00e 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseTaskDataObject.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseTaskDataObject.cs
@@ -256,7 +256,12 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 
         public override void Execute(TaskExecutionMode mode)
         {
+            RestorePlanToExecute = null;
             UpdateRestoreTaskObject();
+            if (IsValid && RestorePlan.RestoreOperations != null && RestorePlan.RestoreOperations.Any())
+            {
+                GetRestorePlanForExecutionAndScript();
+            }
 
             base.Execute(mode);
         }
@@ -291,19 +296,17 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         {
             try
             {
-                if (IsValid && RestorePlan.RestoreOperations != null && RestorePlan.RestoreOperations.Any())
+                if (RestorePlanToExecute != null)
                 {
                     // Restore Plan should be already created and updated at this point
 
-                    RestorePlan restorePlan = GetRestorePlanForExecutionAndScript();
-
-                    if (restorePlan != null && restorePlan.RestoreOperations.Count > 0)
+                    if (RestorePlanToExecute != null && RestorePlanToExecute.RestoreOperations.Count > 0)
                     {
-                        restorePlan.PercentComplete += (object sender, PercentCompleteEventArgs e) =>
+                        RestorePlanToExecute.PercentComplete += (object sender, PercentCompleteEventArgs e) =>
                         {
                             OnMessageAdded(new TaskMessage { Description = $"{e.Percent}%", Status = SqlTaskStatus.InProgress });
                         };
-                        restorePlan.Execute();
+                        RestorePlanToExecute.Execute();
                     }
                 }
                 else