filtering backup sets (#425)

* filtering backup sets

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

@@ -612,7 +612,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                     backupComponent = RestoreConstants.ComponentDatabase;
                     break;
                 case BackupSetType.Differential:
-                    backupType = RestoreConstants.TypeTransactionLog;
+                    backupType = RestoreConstants.TypeDifferential;
                     backupComponent = RestoreConstants.ComponentDatabase;
                     break;                
                 case BackupSetType.FileOrFileGroup:

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

@@ -0,0 +1,75 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.SqlServer.Management.Smo;
+
+namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
+{
+    /// <summary>
+    /// Class include info about selected back sets
+    /// </summary>
+    public class BackupSetsFilterInfo
+    {
+        private HashSet<Guid> selectedBackupSets = new HashSet<Guid>();
+
+        /// <summary>
+        /// Returns true if given backup set is selected
+        /// </summary>
+        public bool IsBackupSetSelected(Guid backupGuid)
+        {
+            bool isSelected = false;
+            if (backupGuid != Guid.Empty)
+            {
+                isSelected = this.selectedBackupSets.Any(x => x == backupGuid);
+            }
+            return isSelected;
+        }
+
+        /// <summary>
+        /// Returns true if given backup set is selected
+        /// </summary>
+        public bool IsBackupSetSelected(BackupSet backupSet)
+        {
+            return IsBackupSetSelected(backupSet != null ? backupSet.BackupSetGuid : Guid.Empty);
+        }
+
+        /// <summary>
+        /// Returns true if any backup set is selected
+        /// </summary>
+        public bool AnySelected
+        {
+            get
+            {
+                return this.selectedBackupSets != null && this.selectedBackupSets.Any();
+            }
+        }
+
+        /// <summary>
+        /// Adds backup set to selected list if not added aleady 
+        /// </summary>
+        /// <param name="backupSet"></param>
+        public void Add(BackupSet backupSet)
+        {
+            if (backupSet != null)
+            {
+                if (!this.selectedBackupSets.Contains(backupSet.BackupSetGuid))
+                {
+                    this.selectedBackupSets.Add(backupSet.BackupSetGuid);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Clears the list
+        /// </summary>
+        public void Clear()
+        {
+            this.selectedBackupSets.Clear();
+        }
+    }
+}

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

@@ -181,7 +181,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 
                         response.PlanDetails.Add(LastBackupTaken, restoreDataObject.GetLastBackupTaken());
 
-                        response.BackupSetsToRestore = restoreDataObject.GetBackupSetInfo().Select(x => new DatabaseFileInfo(x.ConvertPropertiesToArray())).ToArray();
+                        response.BackupSetsToRestore = restoreDataObject.GetSelectedBakupSets();
                         var dbNames = restoreDataObject.GetSourceDbNames();
                         response.DatabaseNamesFromBackupSets = dbNames == null ? new string[] { } : dbNames.ToArray();
                         
@@ -236,14 +236,8 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// </summary>
         private static bool CanRestore(RestoreDatabaseTaskDataObject restoreDataObject)
         {
-            if (restoreDataObject != null)
-            {
-                var backupTypes = restoreDataObject.GetBackupSetInfo();
-                var selectedBackupSets = restoreDataObject.RestoreParams.SelectedBackupSets;
-                return backupTypes.Any(x => (selectedBackupSets == null || selectedBackupSets.Contains(x.GetPropertyValueAsString(DatabaseFileInfo.IdPropertyName))) 
-                && x.BackupType.StartsWith(RestoreConstants.TypeFull));
-            }
-            return false;
+            return restoreDataObject != null && restoreDataObject.RestorePlan != null && restoreDataObject.RestorePlan.RestoreOperations != null
+                && restoreDataObject.RestorePlan.RestoreOperations.Count > 0;
         }
 
         /// <summary>
@@ -266,6 +260,10 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                 this.restoreSessions.AddOrUpdate(sessionId, restoreTaskObject, (key, oldSession) => restoreTaskObject);
                 restoreTaskObject.SessionId = sessionId;
             }
+            else
+            {
+                restoreTaskObject.RestoreParams = restoreParams;
+            }
             return restoreTaskObject;
         }
 
@@ -311,49 +309,46 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// <returns></returns>
         private void UpdateRestorePlan(RestoreDatabaseTaskDataObject restoreDataObject)
         {
-            if (restoreDataObject.PlanUpdateRequired)
+            if (!string.IsNullOrEmpty(restoreDataObject.RestoreParams.BackupFilePaths))
             {
-                if (!string.IsNullOrEmpty(restoreDataObject.RestoreParams.BackupFilePaths))
-                {
-                    restoreDataObject.AddFiles(restoreDataObject.RestoreParams.BackupFilePaths);
-                }
-                restoreDataObject.RestorePlanner.ReadHeaderFromMedia = !string.IsNullOrEmpty(restoreDataObject.RestoreParams.BackupFilePaths);
-
-                if (string.IsNullOrWhiteSpace(restoreDataObject.RestoreParams.SourceDatabaseName))
-                {
-                    restoreDataObject.RestorePlanner.DatabaseName = restoreDataObject.DefaultDbName;
-                }
-                else
-                {
-                    restoreDataObject.RestorePlanner.DatabaseName = restoreDataObject.RestoreParams.SourceDatabaseName;
-                }
-                restoreDataObject.TargetDatabase = restoreDataObject.RestoreParams.TargetDatabaseName;
-
-                restoreDataObject.RestoreOptions.KeepReplication = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.KeepReplication);
-                restoreDataObject.RestoreOptions.ReplaceDatabase = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.ReplaceDatabase);
-                restoreDataObject.RestoreOptions.SetRestrictedUser = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.SetRestrictedUser);
-                string recoveryState = restoreDataObject.RestoreParams.GetOptionValue<string>(RestoreOptionsHelper.RecoveryState);
-                object databaseRecoveryState;
-                if (Enum.TryParse(typeof(DatabaseRecoveryState), recoveryState, out databaseRecoveryState))
-                {
-                    restoreDataObject.RestoreOptions.RecoveryState = (DatabaseRecoveryState)databaseRecoveryState;
-                }
-                bool isTailLogBackupPossible = restoreDataObject.IsTailLogBackupPossible(restoreDataObject.RestorePlanner.DatabaseName);
-                if (isTailLogBackupPossible)
-                {
-                    restoreDataObject.RestorePlanner.BackupTailLog = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.BackupTailLog);
-                    restoreDataObject.TailLogBackupFile = restoreDataObject.RestoreParams.GetOptionValue<string>(RestoreOptionsHelper.TailLogBackupFile);
-                    restoreDataObject.TailLogWithNoRecovery = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.TailLogWithNoRecovery);
-                }
-                else
-                {
-                    restoreDataObject.RestorePlanner.BackupTailLog = false;
-                }
-                
-                restoreDataObject.CloseExistingConnections = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.CloseExistingConnections);
-
-                restoreDataObject.UpdateRestorePlan(restoreDataObject.RestoreParams.RelocateDbFiles);
+                restoreDataObject.AddFiles(restoreDataObject.RestoreParams.BackupFilePaths);
             }
+            restoreDataObject.RestorePlanner.ReadHeaderFromMedia = !string.IsNullOrEmpty(restoreDataObject.RestoreParams.BackupFilePaths);
+
+            if (string.IsNullOrWhiteSpace(restoreDataObject.RestoreParams.SourceDatabaseName))
+            {
+                restoreDataObject.RestorePlanner.DatabaseName = restoreDataObject.DefaultDbName;
+            }
+            else
+            {
+                restoreDataObject.RestorePlanner.DatabaseName = restoreDataObject.RestoreParams.SourceDatabaseName;
+            }
+            restoreDataObject.TargetDatabase = restoreDataObject.RestoreParams.TargetDatabaseName;
+
+            restoreDataObject.RestoreOptions.KeepReplication = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.KeepReplication);
+            restoreDataObject.RestoreOptions.ReplaceDatabase = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.ReplaceDatabase);
+            restoreDataObject.RestoreOptions.SetRestrictedUser = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.SetRestrictedUser);
+            string recoveryState = restoreDataObject.RestoreParams.GetOptionValue<string>(RestoreOptionsHelper.RecoveryState);
+            object databaseRecoveryState;
+            if (Enum.TryParse(typeof(DatabaseRecoveryState), recoveryState, out databaseRecoveryState))
+            {
+                restoreDataObject.RestoreOptions.RecoveryState = (DatabaseRecoveryState)databaseRecoveryState;
+            }
+            bool isTailLogBackupPossible = restoreDataObject.IsTailLogBackupPossible(restoreDataObject.RestorePlanner.DatabaseName);
+            if (isTailLogBackupPossible)
+            {
+                restoreDataObject.RestorePlanner.BackupTailLog = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.BackupTailLog);
+                restoreDataObject.TailLogBackupFile = restoreDataObject.RestoreParams.GetOptionValue<string>(RestoreOptionsHelper.TailLogBackupFile);
+                restoreDataObject.TailLogWithNoRecovery = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.TailLogWithNoRecovery);
+            }
+            else
+            {
+                restoreDataObject.RestorePlanner.BackupTailLog = false;
+            }
+
+            restoreDataObject.CloseExistingConnections = restoreDataObject.RestoreParams.GetOptionValue<bool>(RestoreOptionsHelper.CloseExistingConnections);
+
+            restoreDataObject.UpdateRestorePlan(restoreDataObject.RestoreParams.RelocateDbFiles);
         }
 
         /// <summary>

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

@@ -23,6 +23,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         private const char BackupMediaNameSeparator = ',';
         private DatabaseRestorePlanner restorePlanner;
         private string tailLogBackupFile;
+        private BackupSetsFilterInfo backupSetsFilterInfo = new BackupSetsFilterInfo();
 
         public RestoreDatabaseTaskDataObject(Server server, String databaseName)
         {
@@ -132,25 +133,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             }
         }
 
-        /// <summary>
-        /// Removes the backup sets that are filtered in the request
-        /// </summary>
-        public void RemoveFilteredBackupSets()
-        {
-            var backupSetIdsToRestore = RestoreParams.SelectedBackupSets;
-            if (backupSetIdsToRestore != null)
-            {
-                var ids = backupSetIdsToRestore.Select(x =>
-                {
-                    Guid guid;
-                    Guid.TryParse(x, out guid);
-                    return guid;
-                }
-                );
-                restorePlan.RestoreOperations.RemoveAll(x => !ids.Contains(x.BackupSet.BackupSetGuid));
-            }
-        }
-
+        
         /// <summary>
         /// Returns the last backup taken
         /// </summary>
@@ -158,20 +141,22 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         public string GetLastBackupTaken()
         {
             string lastBackup = string.Empty;
-            int lastIndexSel = 0; //TODO: find the selected backup set
+            int lastSelectedIndex = GetLastSelectedBackupSetIndex();
+            BackupSet lastSelectedBackupSet = lastSelectedIndex >= 0 && this.RestorePlan.RestoreOperations != null && this.RestorePlan.RestoreOperations.Count > 0 ? 
+                this.RestorePlan.RestoreOperations[lastSelectedIndex].BackupSet : null;
             if (this.RestorePlanner.RestoreToLastBackup &&
-                this.RestorePlan.RestoreOperations[lastIndexSel] != null &&
+                lastSelectedBackupSet != null &&
                 this.RestorePlan.RestoreOperations.Count > 0 &&
-                this.RestorePlan.RestoreOperations[lastIndexSel].BackupSet != null)
+                lastSelectedBackupSet != null)
             {
-                int lastIndex = this.RestorePlan.RestoreOperations.Count - 1;
-                DateTime backupTime = this.RestorePlan.RestoreOperations[lastIndexSel].BackupSet.BackupStartDate;
+                bool isTheLastOneSelected = lastSelectedIndex == this.RestorePlan.RestoreOperations.Count - 1;
+                DateTime backupTime = lastSelectedBackupSet.BackupStartDate;
                 string backupTimeStr = backupTime.ToLongDateString() + " " + backupTime.ToLongTimeString();
-                lastBackup = (lastIndexSel == lastIndex) ?
+                lastBackup = isTheLastOneSelected ?
                     string.Format(CultureInfo.CurrentCulture, SR.TheLastBackupTaken, (backupTimeStr)) : backupTimeStr;
             }
             //TODO: find the selected one
-            else if (this.RestoreSelected[0] && !this.RestorePlanner.RestoreToLastBackup)
+            else if (GetFirstSelectedBackupSetIndex() == 0 && !this.RestorePlanner.RestoreToLastBackup)
             {
                 lastBackup = this.CurrentRestorePointInTime.Value.ToLongDateString() +
                         " " + this.CurrentRestorePointInTime.Value.ToLongTimeString();
@@ -185,13 +170,12 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// </summary>
         public void Execute()
         {
-            RestorePlan restorePlan = RestorePlan;
+            RestorePlan restorePlan = GetRestorePlanForExecutionAndScript();
             // ssms creates a new restore plan by calling GetRestorePlanForExecutionAndScript and
             // Doens't use the plan already created here. not sure why, using the existing restore plan doesn't make
             // any issue so far so keeping in it for now but we might want to double check later
             if (restorePlan != null && restorePlan.RestoreOperations.Count > 0)
             {
-                RemoveFilteredBackupSets();
                 restorePlan.PercentComplete += (object sender, PercentCompleteEventArgs e) =>
                 {
                     if (SqlTask != null)
@@ -203,6 +187,49 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             }
         }
 
+        /// <summary>
+        /// Gets RestorePlan to perform restore and to script
+        /// </summary>
+        public RestorePlan GetRestorePlanForExecutionAndScript()
+        {
+            this.ActiveException = null; //Clear any existing exceptions as the plan is getting recreated. 
+            //Clear any existing exceptions as new plan is getting recreated.
+            this.CreateOrUpdateRestorePlanException = null;
+            bool tailLogBackup = this.RestorePlanner.BackupTailLog;
+            if (this.PlanUpdateRequired)
+            {
+                this.RestorePlan = this.RestorePlanner.CreateRestorePlan(this.RestoreOptions);
+                
+                this.Util.AddCredentialNameForUrlBackupSet(this.RestorePlan, this.CredentialName);
+            }
+            RestorePlan rp = new RestorePlan(this.Server);
+            rp.RestoreAction = RestoreActionType.Database;
+            if (this.RestorePlan != null)
+            {
+                if (this.RestorePlan.TailLogBackupOperation != null && tailLogBackup)
+                {
+                    rp.TailLogBackupOperation = this.RestorePlan.TailLogBackupOperation;
+                }
+                int i = 0;
+                foreach (Restore res in this.RestorePlan.RestoreOperations)
+                {
+                    if (this.backupSetsFilterInfo.IsBackupSetSelected(res.BackupSet))
+                    {
+                        rp.RestoreOperations.Add(res);
+                    }
+                    i++;
+                }
+            }
+            this.SetRestorePlanProperties(rp);
+            RestorePlanToExecute = rp;
+            return rp;
+        }
+
+        /// <summary>
+        /// For test purpose only. The restore plan that's used to execute
+        /// </summary>
+        internal RestorePlan RestorePlanToExecute { get; set; }
+
         /// <summary>
         /// Restore Util
         /// </summary>
@@ -478,8 +505,6 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             }
         }
 
-        public bool[] RestoreSelected;
-
         /// <summary>
         /// The database being restored
         /// </summary>
@@ -526,34 +551,6 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// </summary>
         internal string ContainerSharedAccessPolicy = string.Empty;
 
-        /// <summary>
-        /// Updates the RestoreSelected Array to hold information about updated Restore Plan
-        /// </summary>
-        private void UpdateRestoreSelected()
-        {
-            int operationsCount = this.RestorePlan.RestoreOperations.Count;
-            // The given condition will return true only if new backup has been added on database during lifetime of restore dialog.
-            // This will happen when tail log backup is taken successfully and subsequent restores have failed.
-            if (operationsCount > this.RestoreSelected.Length)
-            {
-                bool[] tempRestoreSel = new bool[this.RestorePlan.RestoreOperations.Count];
-                for (int i = 0; i < operationsCount; i++)
-                {
-                    if (i < RestoreSelected.Length)
-                    {
-                        //Retain all the old values.
-                        tempRestoreSel[i] = RestoreSelected[i];
-                    }
-                    else
-                    {
-                        //Do not add the newly added backupset into Restore plan by default.
-                        tempRestoreSel[i] = false;
-                    }
-                }
-                this.RestoreSelected = tempRestoreSel;
-            }
-        }
-
         /// <summary>
         /// Returns the physical name for the target Db file.
         /// It is the sourceDbName replaced with targetDbName in sourceFilename.
@@ -577,12 +574,38 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             List<BackupSetInfo> result = new List<BackupSetInfo>();
             foreach (Restore restore in RestorePlan.RestoreOperations)
             {
-                result.Add(BackupSetInfo.Create(restore, Server));
+                BackupSetInfo backupSetInfo = BackupSetInfo.Create(restore, Server);
+                if (this.backupSetsFilterInfo.IsBackupSetSelected(restore.BackupSet))
+                {
+                   
+                }
+                result.Add(backupSetInfo);
             }
 
             return result;
         }
 
+        /// <summary>
+        /// List of selected backupsets
+        /// </summary>
+        public DatabaseFileInfo[] GetSelectedBakupSets()
+        {
+            List<DatabaseFileInfo> result = new List<DatabaseFileInfo>();
+            IEnumerable<BackupSetInfo> backupSetInfos = GetBackupSetInfo();
+            foreach (var backupSetInfo in backupSetInfos)
+            {
+                var item = new DatabaseFileInfo(backupSetInfo.ConvertPropertiesToArray());
+                Guid backupSetGuid;
+                if(!Guid.TryParse(item.Id, out backupSetGuid))
+                {
+                    backupSetGuid = Guid.Empty;
+                }
+                item.IsSelected = this.backupSetsFilterInfo.IsBackupSetSelected(backupSetGuid);
+                result.Add(item);
+            }
+            return result.ToArray();
+        }
+
         /// <summary>
         /// Gets the files of the database
         /// </summary>
@@ -682,23 +705,15 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                     this.SetRestorePlanProperties(this.restorePlan);
                 }
             }
-            if (this.restorePlan != null)
-            {
-                this.RestoreSelected = new bool[this.restorePlan.RestoreOperations.Count];
-                for (int i = 0; i < this.restorePlan.RestoreOperations.Count; i++)
-                {
-                    this.RestoreSelected[i] = true;
-                }
-            }
-            else
+            if (this.restorePlan == null)
             {
                 this.RestorePlan = new RestorePlan(this.Server);
                 this.Util.AddCredentialNameForUrlBackupSet(this.RestorePlan, this.CredentialName);
-                this.RestoreSelected = new bool[0];
             }
+
+            UpdateSelectedBackupSets();
         }
 
-        
         /// <summary>
         /// Determine if restore plan of selected database does have Url
         /// </summary>
@@ -720,9 +735,9 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             }
             rp.SetRestoreOptions(this.RestoreOptions);
             rp.CloseExistingConnections = this.CloseExistingConnections;
-            if (this.targetDbName != null && !this.targetDbName.Equals(string.Empty))
+            if (this.TargetDatabase != null && !this.TargetDatabase.Equals(string.Empty))
             {
-                rp.DatabaseName = targetDbName;
+                rp.DatabaseName = TargetDatabase;
             }
             rp.RestoreOperations[0].RelocateFiles.Clear();
             foreach (DbFile dbFile in this.DbFiles)
@@ -789,109 +804,150 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         {
             get
             {
-                if (this.RestorePlan == null || this.RestorePlan.RestoreOperations.Count == 0
-                    || this.RestoreSelected.Length == 0 || !this.RestoreSelected[0])
+                if (!IsAnyFullBackupSetSelected())
                 {
                     return null;
                 }
                 for (int i = this.RestorePlan.RestoreOperations.Count - 1; i >= 0; i--)
                 {
-                    if (this.RestoreSelected[i])
+                    BackupSet backupSet = this.RestorePlan.RestoreOperations[i].BackupSet;
+                    if (this.backupSetsFilterInfo.IsBackupSetSelected(backupSet))
                     {
-                        if (this.RestorePlan.RestoreOperations[i].BackupSet == null
-                            || (this.RestorePlan.RestoreOperations[i].BackupSet.BackupSetType == BackupSetType.Log &&
+                        if (backupSet == null
+                            || (backupSet.BackupSetType == BackupSetType.Log &&
                                 this.RestorePlan.RestoreOperations[i].ToPointInTime != null))
                         {
                             return this.RestorePlanner.RestoreToPointInTime;
                         }
-                        return this.RestorePlan.RestoreOperations[i].BackupSet.BackupStartDate;
+                        return backupSet.BackupStartDate;
                     }
                 }
                 return null;
             }
         }
 
-        public void ToggleSelectRestore(int index)
+
+
+       
+
+        private bool IsAnyFullBackupSetSelected()
         {
-            RestorePlan rp = this.restorePlan;
-            if (rp == null || rp.RestoreOperations.Count <= index)
+            bool isSelected = false;
+
+            if (this.RestorePlan != null && this.RestorePlan.RestoreOperations.Any() && this.backupSetsFilterInfo.AnySelected)
             {
-                return;
+                var fullBackupSet = this.RestorePlan.RestoreOperations.FirstOrDefault(x => x.BackupSet.BackupSetType == BackupSetType.Database);
+                isSelected = fullBackupSet != null && this.backupSetsFilterInfo.IsBackupSetSelected(fullBackupSet.BackupSet.BackupSetGuid);
             }
-            //the last index - this will include tail-Log restore operation if present
-            if (index == rp.RestoreOperations.Count - 1)
+
+            return isSelected;
+        }
+
+        private int GetLastSelectedBackupSetIndex()
+        {
+            if (this.RestorePlan != null && this.RestorePlan.RestoreOperations.Any() && this.backupSetsFilterInfo.AnySelected)
             {
-                if (this.RestoreSelected[index])
+                for (int i = this.RestorePlan.RestoreOperations.Count -1; i >= 0; i--)
                 {
-                    this.RestoreSelected[index] = false;
-                }
-                else
-                {
-                    for (int i = 0; i <= index; i++)
+                    BackupSet backupSet = this.RestorePlan.RestoreOperations[i].BackupSet;
+                    if (this.backupSetsFilterInfo.IsBackupSetSelected(backupSet))
                     {
-                        this.RestoreSelected[i] = true;
+                        return i;
                     }
                 }
-                return;
             }
-            if (index == 0)
+            return -1;
+        }
+
+        private int GetFirstSelectedBackupSetIndex()
+        {
+            if (this.RestorePlan != null && this.RestorePlan.RestoreOperations.Any() && this.backupSetsFilterInfo.AnySelected)
             {
-                if (!this.RestoreSelected[index])
+                for (int i = 0; i < this.RestorePlan.RestoreOperations.Count - 1; i++)
                 {
-                    this.RestoreSelected[index] = true;
-                }
-                else
-                {
-                    for (int i = index; i < rp.RestoreOperations.Count; i++)
+                    BackupSet backupSet = this.RestorePlan.RestoreOperations[i].BackupSet;
+                    if (this.backupSetsFilterInfo.IsBackupSetSelected(backupSet))
                     {
-                        this.RestoreSelected[i] = false;
+                        return i;
                     }
                 }
-                return;
             }
-            
-            if (index == 1 && rp.RestoreOperations[index].BackupSet.BackupSetType == BackupSetType.Differential)
+            return -1;
+        }
+
+
+        public void UpdateSelectedBackupSets()
+        {
+            this.backupSetsFilterInfo.Clear();
+            var selectedBackupSetsFromClient = this.RestoreParams.SelectedBackupSets;
+
+            if (this.RestorePlan != null && this.RestorePlan.RestoreOperations != null)
             {
-                if (!this.RestoreSelected[index])
+                for (int index = 0; index < this.RestorePlan.RestoreOperations.Count; index++)
                 {
-                    this.RestoreSelected[0] = true;
-                    this.RestoreSelected[index] = true;
-                }
-                else if (rp.RestoreOperations[2].BackupSet == null)
-                {
-                    this.RestoreSelected[index] = false;
-                    this.RestoreSelected[2] = false;
-                }
-                else if (this.Server.Version.Major < 9 || BackupSet.IsBackupSetsInSequence(rp.RestoreOperations[0].BackupSet, rp.RestoreOperations[2].BackupSet))
-                {
-                    this.RestoreSelected[index] = false;
-                }
-                else
-                {
-                    for (int i = index; i < rp.RestoreOperations.Count; i++)
+                    BackupSet backupSet = this.RestorePlan.RestoreOperations[index].BackupSet;
+                    if (backupSet != null)
                     {
-                        this.RestoreSelected[i] = false;
+                        // If the collection client sent is null, select everything; otherwise select the items that are selected in client
+                        bool backupSetSelected = selectedBackupSetsFromClient == null || selectedBackupSetsFromClient.Any(x => BackUpSetGuidEqualsId(backupSet, x));
+
+                        if (backupSetSelected)
+                        {
+                            AddBackupSetsToSelected(index, index);
+
+                            //the last index - this will include tail-Log restore operation if present
+                            //If the last item is selected, select all the items before that because the last item
+                            //is tail-log
+                            if (index == this.RestorePlan.RestoreOperations.Count - 1)
+                            {
+                                AddBackupSetsToSelected(0, index);
+                            }
+
+                            //If the second item is selected and it's a diff backup, the fist item (full backup) has to be selected
+                            if (index == 1 && backupSet.BackupSetType == BackupSetType.Differential)
+                            {
+                                AddBackupSetsToSelected(0, 0);
+                            }
+
+                            //If the selected item is a log backup, select all the items before that
+                            if (backupSet.BackupSetType == BackupSetType.Log)
+                            {
+                                AddBackupSetsToSelected(0, index);
+                            }
+                        }
+                        else
+                        {
+                            // If the first item is not selected which is always the full backup, other backupsets cannot be selected
+                            if (index == 0)
+                            {
+                                selectedBackupSetsFromClient = new string[] { };
+                            }
+
+                            //The a log is not selected, the logs after that cannot be selected
+                            if (backupSet.BackupSetType == BackupSetType.Log)
+                            {
+                                break;
+                            }
+                        }
                     }
                 }
-                return;
             }
-            if (rp.RestoreOperations[index].BackupSet.BackupSetType == BackupSetType.Log)
+        }
+
+        private bool BackUpSetGuidEqualsId(BackupSet backupSet, string id)
+        {
+           return backupSet != null && string.Compare(backupSet.BackupSetGuid.ToString(), id, StringComparison.OrdinalIgnoreCase) == 0;
+        }
+
+        private void AddBackupSetsToSelected(int from, int to)
+        {
+            if (this.RestorePlan != null && this.RestorePlan.RestoreOperations != null 
+                && from < this.RestorePlan.RestoreOperations.Count && to < this.RestorePlan.RestoreOperations.Count)
             {
-                if (this.RestoreSelected[index])
+                for (int i = from; i <= to; i++)
                 {
-                    for (int i = index; i < rp.RestoreOperations.Count; i++)
-                    {
-                        this.RestoreSelected[i] = false;
-                    }
-                    return;
-                }
-                else
-                {
-                    for (int i = 0; i <= index; i++)
-                    {
-                        this.RestoreSelected[i] = true;
-                    }
-                    return;
+                    BackupSet backupSet = this.RestorePlan.RestoreOperations[i].BackupSet;
+                    this.backupSetsFilterInfo.Add(backupSet);
                 }
             }
         }

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

@@ -12,7 +12,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
         internal const string KeepReplication = "keepReplication";
         internal const string ReplaceDatabase = "replaceDatabase";
         internal const string SetRestrictedUser = "setRestrictedUser";
-        internal const string RecoveryState = "eecoveryState";
+        internal const string RecoveryState = "recoveryState";
         internal const string BackupTailLog = "backupTailLog";
         internal const string DefaultBackupTailLog = "defaultBackupTailLog";
         internal const string TailLogBackupFile = "tailLogBackupFile";
@@ -48,7 +48,8 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                     Description = "Preserve the replication settings (WITH KEEP_REPLICATION)",
                     ValueType = ServiceOption.ValueTypeBoolean,
                     IsRequired = false,
-                    GroupName = "Restore options"
+                    GroupName = "Restore options",
+                    DefaultValue = "false"
                 },
                 new ServiceOption
                 {
@@ -57,7 +58,8 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                     Description = "Overwrite the existing database (WITH REPLACE)",
                     ValueType = ServiceOption.ValueTypeBoolean,
                     IsRequired = false,
-                    GroupName = "Restore options"
+                    GroupName = "Restore options",
+                    DefaultValue = "false"
                 },
                 new ServiceOption
                 {
@@ -66,7 +68,8 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                     Description = "Restrict access to the restored database (WITH RESTRICTED_USER)",
                     ValueType = ServiceOption.ValueTypeBoolean,
                     IsRequired = false,
-                    GroupName = "Restore options"
+                    GroupName = "Restore options",
+                    DefaultValue = "false"
                 },
                 new ServiceOption
                 {
@@ -93,7 +96,8 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                             Name = "WithStandBy",
                             DisplayName = "RESTORE WITH STANDBY"
                         }
-                    }
+                    },
+                    DefaultValue = "WithRecovery"
                 },
                 new ServiceOption
                 {
@@ -158,7 +162,8 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                     Description = "Relocate all files",
                     ValueType = ServiceOption.ValueTypeBoolean,
                     IsRequired = false,
-                    GroupName = "Restore database files as"
+                    GroupName = "Restore database files as",
+                    DefaultValue = "false"
                 },
                 new ServiceOption
                 {

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

@@ -3,6 +3,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
 using System.Collections.Generic;
 using System.Data.SqlClient;
 using System.Globalization;
@@ -32,7 +33,12 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
         private ConnectionService _connectService = TestServiceProvider.Instance.ConnectionService;
         private Mock<IProtocolEndpoint> serviceHostMock;
         private DisasterRecoveryService service;
-        private string fullBackUpDatabase;
+        private string fullBackupFilePath;
+        private string[] backupFilesToRecoverDatabase;
+
+        //The table names used in the script to create backup files for a database
+        //Each table is created after a backup script to verify recovering to different states
+        private string[] tableNames = new string[] { "tb1", "tb2", "tb3", "tb4", "tb5" };
 
         public RestoreDatabaseServiceTests()
         {
@@ -43,18 +49,110 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
 
         private async Task VerifyBackupFileCreated()
         {
-            if(fullBackUpDatabase == null)
+            if(fullBackupFilePath == null)
             {
-                fullBackUpDatabase = await CreateBackupFile();
+                fullBackupFilePath = await CreateBackupFile();
             }
         }
 
+        private async Task<string[]> GetBackupFilesToRecoverDatabaseCreated()
+        {
+            if(backupFilesToRecoverDatabase == null)
+            {
+                backupFilesToRecoverDatabase = await CreateBackupSetsToRecoverDatabase();
+            }
+            return backupFilesToRecoverDatabase;
+        }
+
         [Fact]
         public async void RestorePlanShouldCreatedSuccessfullyForFullBackup()
         {
             await VerifyBackupFileCreated();
             bool canRestore = true;
-            await VerifyRestore(fullBackUpDatabase, canRestore);
+            await VerifyRestore(fullBackupFilePath, canRestore);
+        }
+
+        [Fact]
+        public async void RestoreShouldNotRestoreAnyBackupSetsIfFullNotSelected()
+        {
+            var backupFiles = await GetBackupFilesToRecoverDatabaseCreated();
+            //Remove the full backupset
+            int indexToDelete = 0;
+            //Verify that all backupsets are restored
+            int[] expectedTable = new int[] { };
+
+            await VerifyRestoreMultipleBackupSets(backupFiles, indexToDelete, expectedTable);
+        }
+
+        [Fact]
+        public async void RestoreShouldRestoreTheBackupSetsThatAreSelected()
+        {
+            var backupFiles = await GetBackupFilesToRecoverDatabaseCreated();
+            //Remove the last backupset
+            int indexToDelete = 4;
+            //Verify that backupset is not restored
+            int[] expectedTable = new int[] { 0, 1, 2, 3 };
+
+            await VerifyRestoreMultipleBackupSets(backupFiles, indexToDelete, expectedTable);
+        }
+
+        [Fact]
+        public async void RestoreShouldNotRestoreTheLogBackupSetsIfOneNotSelected()
+        {
+            var backupFiles = await GetBackupFilesToRecoverDatabaseCreated();
+            //Remove the one of the log backup sets
+            int indexToDelete = 3;
+            //Verify the logs backup set that's removed and all logs after that are not restored
+            int[] expectedTable = new int[] { 0, 1, 2 };
+            await VerifyRestoreMultipleBackupSets(backupFiles, indexToDelete, expectedTable);
+        }
+
+        private async Task VerifyRestoreMultipleBackupSets(string[] backupFiles, int backupSetIndexToDelete, int[] expectedSelectedIndexes)
+        {
+            var testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, null, "RestoreTest");
+            try
+            {
+                string targetDbName = testDb.DatabaseName;
+                bool canRestore = true;
+                var response = await VerifyRestore(backupFiles, canRestore, false, targetDbName, null, null);
+                Assert.True(response.BackupSetsToRestore.Count() >= 2);
+                var allIds = response.BackupSetsToRestore.Select(x => x.Id).ToList();
+                if (backupSetIndexToDelete >= 0)
+                {
+                    allIds.RemoveAt(backupSetIndexToDelete);
+                }
+                string[] selectedIds = allIds.ToArray();
+                Dictionary<string, object> options = new Dictionary<string, object>();
+                options.Add(RestoreOptionsHelper.ReplaceDatabase, true);
+                response = await VerifyRestore(backupFiles, canRestore, true, targetDbName, selectedIds, options, (database) =>
+                {
+                    bool tablesFound = true;
+                    for (int i = 0; i < tableNames.Length; i++)
+                    {
+                        string tableName = tableNames[i];
+                        if (!database.Tables.Contains(tableName) && expectedSelectedIndexes.Contains(i))
+                        {
+                            tablesFound = false;
+                            break;
+                        }
+                    }
+                    bool numberOfTableCreatedIsCorrect = database.Tables.Count == expectedSelectedIndexes.Length;
+                    return numberOfTableCreatedIsCorrect && tablesFound;
+                });
+
+                for (int i = 0; i < response.BackupSetsToRestore.Count(); i++)
+                {
+                    DatabaseFileInfo databaseInfo = response.BackupSetsToRestore[i];
+                    Assert.Equal(databaseInfo.IsSelected, expectedSelectedIndexes.Contains(i));
+                }
+            }
+            finally
+            {
+                if (testDb != null)
+                {
+                   testDb.Cleanup();
+                }
+            }
         }
 
         [Fact]
@@ -70,7 +168,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                 Dictionary<string, object> options = new Dictionary<string, object>();
                 options.Add(RestoreOptionsHelper.ReplaceDatabase, true);
 
-                await VerifyRestore(new string[] { fullBackUpDatabase }, canRestore, true, testDb.DatabaseName, null, options);
+                await VerifyRestore(new string[] { fullBackupFilePath }, canRestore, true, testDb.DatabaseName, null, options);
             }
             finally
             {
@@ -92,7 +190,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                 await VerifyBackupFileCreated();
                 bool canRestore = true;
 
-                await VerifyRestore(new string[] { fullBackUpDatabase }, canRestore, false, testDb.DatabaseName, null, null);
+                await VerifyRestore(new string[] { fullBackupFilePath }, canRestore, false, testDb.DatabaseName, null, null);
             }
             finally
             {
@@ -125,12 +223,12 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
             if(fileInfo != null)
             {
                 var selectedBackupSets = new string[] { fileInfo.Id };
-                await VerifyRestore(backupFileNames, false, false, "RestoredFromTwoBackupFile", selectedBackupSets);
+                await VerifyRestore(backupFileNames, true, false, "RestoredFromTwoBackupFile", selectedBackupSets);
             }
         }
 
         [Fact]
-        public async void RestoreShouldCompletedSuccessfullyGivenTowBackupFilesButFilterDifferentialBackup()
+        public async void RestoreShouldCompletedSuccessfullyGivenTwoBackupFilesButFilterDifferentialBackup()
         {
 
             string[] backupFileNames = new string[] { "FullBackup.bak", "DiffBackup.bak" };
@@ -150,7 +248,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
         {
             await VerifyBackupFileCreated();
 
-            string backupFileName = fullBackUpDatabase;
+            string backupFileName = fullBackupFilePath;
             bool canRestore = true;
             var restorePlan = await VerifyRestore(backupFileName, canRestore, true);
             Assert.NotNull(restorePlan.BackupSetsToRestore);
@@ -161,24 +259,24 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
         {
             await VerifyBackupFileCreated();
 
-            string backupFileName = fullBackUpDatabase;
+            string backupFileName = fullBackupFilePath;
             bool canRestore = true;
             var restorePlan = await VerifyRestore(backupFileName, canRestore, true, "NewRestoredDatabase");
         }
 
         [Fact]
-        public async void RestorePlanShouldFailForDiffBackup()
+        public async void RestorePlanShouldCreatedSuccessfullyForDiffBackup()
         {
             string backupFileName = "DiffBackup.bak";
-            bool canRestore = false;
+            bool canRestore = true;
             await VerifyRestore(backupFileName, canRestore);
         }
 
         [Fact]
-        public async void RestorePlanShouldFailForTransactionLogBackup()
+        public async void RestorePlanShouldCreatedSuccessfullyForTransactionLogBackup()
         {
             string backupFileName = "TransactionLogBackup.bak";
-            bool canRestore = false;
+            bool canRestore = true;
             await VerifyRestore(backupFileName, canRestore);
         }
 
@@ -191,7 +289,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
             {
                 TestConnectionResult connectionResult = await LiveConnectionHelper.InitLiveConnectionInfoAsync("master", queryTempFile.FilePath);
 
-                string filePath = GetBackupFilePath(fullBackUpDatabase);
+                string filePath = GetBackupFilePath(fullBackupFilePath);
 
                 RestoreParams restoreParams = new RestoreParams
                 {
@@ -279,7 +377,8 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
             bool execute = false, 
             string targetDatabase = null, 
             string[] selectedBackupSets = null,
-            Dictionary<string, object> options = null)
+            Dictionary<string, object> options = null,
+            Func<Database, bool> verifyDatabase = null)
         {
             var filePaths = backupFileNames.Select(x => GetBackupFilePath(x));
             string backUpFilePath = filePaths.Aggregate((current, next) => current + " ," + next);
@@ -337,15 +436,22 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                         request.SessionId = response.SessionId;
                         restoreDataObject = service.CreateRestoreDatabaseTaskDataObject(request);
                         Assert.Equal(response.SessionId, restoreDataObject.SessionId);
-                        //await DropDatabase(targetDatabase);
-                        //Thread.Sleep(2000);
                         request.RelocateDbFiles = !restoreDataObject.DbFilesLocationAreValid();
                         service.ExecuteRestore(restoreDataObject);
                         Assert.True(restoreDataObject.Server.Databases.Contains(targetDatabase));
-                        if(selectedBackupSets != null)
+
+                        if(verifyDatabase != null)
                         {
-                            Assert.Equal(selectedBackupSets.Count(), restoreDataObject.RestorePlan.RestoreOperations.Count());
+                            Assert.True(verifyDatabase(restoreDataObject.Server.Databases[targetDatabase]));
                         }
+
+                        //To verify the backupset that are restored, verifying the database is a better options.
+                        //Some tests still verify the number of backup sets that are executed which in some cases can be less than the selected list
+                        if (verifyDatabase == null && selectedBackupSets != null)
+                        {
+                            Assert.Equal(selectedBackupSets.Count(), restoreDataObject.RestorePlanToExecute.RestoreOperations.Count());
+                        }
+                     
                         await DropDatabase(targetDatabase);
                     }
                 }
@@ -403,6 +509,60 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                .RegisterSingleService(new DisasterRecoveryService());
         }
 
+        public async Task<string[]> CreateBackupSetsToRecoverDatabase()
+        {
+            List<string> backupFiles = new List<string>();
+            using (SelfCleaningTempFile queryTempFile = new SelfCleaningTempFile())
+            {
+                string query = $"create table {tableNames[0]} (c1 int)";
+                SqlTestDb testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, query, "RestoreTest");
+                string databaseName = testDb.DatabaseName;
+                // Initialize backup service
+                var liveConnection = LiveConnectionHelper.InitLiveConnectionInfo(databaseName, queryTempFile.FilePath);
+                DatabaseTaskHelper helper = AdminService.CreateDatabaseTaskHelper(liveConnection.ConnectionInfo, databaseExists: true);
+                SqlConnection sqlConn = DisasterRecoveryService.GetSqlConnection(liveConnection.ConnectionInfo);
+                BackupConfigInfo backupConfigInfo = DisasterRecoveryService.Instance.GetBackupConfigInfo(helper.DataContainer, sqlConn, sqlConn.Database);
+
+                string backupPath = Path.Combine(backupConfigInfo.DefaultBackupFolder, databaseName + "_full.bak");
+                query = $"BACKUP DATABASE [{databaseName}] TO  DISK = N'{backupPath}' WITH NOFORMAT, NOINIT, NAME = N'{databaseName}-Full Database Backup', SKIP, NOREWIND, NOUNLOAD,  STATS = 10";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, "master", query);
+                backupFiles.Add(backupPath);
+
+                query = $"create table {tableNames[1]} (c1 int)";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, databaseName, query);
+                backupPath = Path.Combine(backupConfigInfo.DefaultBackupFolder, databaseName + "_diff.bak");
+                query = $"BACKUP DATABASE [{databaseName}] TO  DISK = N'{backupPath}' WITH DIFFERENTIAL, NOFORMAT, NOINIT, NAME = N'{databaseName}-Full Database Backup', SKIP, NOREWIND, NOUNLOAD,  STATS = 10";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, "master", query);
+                backupFiles.Add(backupPath);
+
+                query = $"create table {tableNames[2]} (c1 int)";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, databaseName, query);
+                backupPath = Path.Combine(backupConfigInfo.DefaultBackupFolder, databaseName + "_log1.bak");
+                query = $"BACKUP Log [{databaseName}] TO  DISK = N'{backupPath}' WITH NOFORMAT, NOINIT, NAME = N'{databaseName}-Full Database Backup', SKIP, NOREWIND, NOUNLOAD,  STATS = 10";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, "master", query);
+                backupFiles.Add(backupPath);
+
+                query = $"create table {tableNames[3]} (c1 int)";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, databaseName, query);
+                backupPath = Path.Combine(backupConfigInfo.DefaultBackupFolder, databaseName + "_log2.bak");
+                query = $"BACKUP Log [{databaseName}] TO  DISK = N'{backupPath}' WITH NOFORMAT, NOINIT, NAME = N'{databaseName}-Full Database Backup', SKIP, NOREWIND, NOUNLOAD,  STATS = 10";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, "master", query);
+                backupFiles.Add(backupPath);
+
+                query = $"create table {tableNames[4]} (c1 int)";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, databaseName, query);
+                backupPath = Path.Combine(backupConfigInfo.DefaultBackupFolder, databaseName + "_log3.bak");
+                query = $"BACKUP Log [{databaseName}] TO  DISK = N'{backupPath}' WITH NOFORMAT, NOINIT, NAME = N'{databaseName}-Full Database Backup', SKIP, NOREWIND, NOUNLOAD,  STATS = 10";
+                await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, "master", query);
+                backupFiles.Add(backupPath);
+
+                // Clean up the database
+                testDb.Cleanup();
+            }
+            return backupFiles.ToArray();
+
+        }
+
         public async Task<string> CreateBackupFile()
         {
             using (SelfCleaningTempFile queryTempFile = new SelfCleaningTempFile())