Adding more features to restore operation (#420)

* Adding more features to restore operations and added tests

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

@@ -0,0 +1,60 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.Linq;
+using Microsoft.SqlTools.Utility;
+
+namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts
+{
+    /// <summary>
+    /// Class includes information about a file related to a database operation. 
+    /// Can be used for backup set files or restored database files
+    /// </summary>
+    public class DatabaseFileInfo
+    {
+        /// <summary>
+        /// The property name used for ids
+        /// </summary>
+        public const string IdPropertyName = "Id";
+
+        public DatabaseFileInfo(LocalizedPropertyInfo[] properties)
+        {
+            Validate.IsNotNull("properties", properties);
+
+            this.Properties = properties;
+            if (this.Properties != null )
+            {
+                var idProperty = this.Properties.FirstOrDefault(x => x.PropertyName == IdPropertyName);
+                Id = idProperty == null || idProperty.PropertyValue == null ? string.Empty : idProperty.PropertyValue.ToString();
+            }
+        }
+
+        /// <summary>
+        /// Properties
+        /// </summary>
+        public LocalizedPropertyInfo[] Properties { get; private set; }
+
+        public string GetPropertyValueAsString(string name)
+        {
+            string value = string.Empty;
+            if (Properties != null)
+            {
+                var property = Properties.FirstOrDefault(x => x.PropertyName == name);
+                value = property == null || property.PropertyValue == null ? string.Empty : property.PropertyValue.ToString();
+            }
+            return value;
+        }
+
+        /// <summary>
+        /// Unique id for this item
+        /// </summary>
+        public string Id { get; private set; }
+
+        /// <summary>
+        /// Indicates whether the item is selected in client
+        /// </summary>
+        public bool IsSelected { get; set; }
+    }
+}

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

@@ -0,0 +1,72 @@
+//
+// 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.Globalization;
+
+namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts
+{
+    public class LocalizedPropertyInfo
+    {
+
+        private string propertyValueDisplayName;
+        private string propertyDisplayName;
+
+        /// <summary>
+        /// Property name
+        /// </summary>
+        public string PropertyName { get; set; }
+
+
+        /// <summary>
+        /// Property value
+        /// </summary>
+        public object PropertyValue { get; set; }
+
+        /// <summary>
+        /// Property display name
+        /// </summary>
+        public string PropertyDisplayName
+        {
+            get
+            {
+                return string.IsNullOrEmpty(this.propertyDisplayName) ? PropertyName : this.propertyDisplayName;
+            }
+            set
+            {
+                this.propertyDisplayName = value;
+            }
+        }
+
+        /// <summary>
+        /// Property display name for the value
+        /// </summary>
+        public string PropertyValueDisplayName
+        {
+            get
+            {
+                return string.IsNullOrEmpty(propertyValueDisplayName) ? GetLocalizedPropertyValue() : propertyValueDisplayName;
+            }
+            set
+            {
+                this.propertyValueDisplayName = value;
+            }
+        }
+
+        private string GetLocalizedPropertyValue()
+        {
+            string displayName = string.Empty;
+            if(PropertyValue is DateTime)
+            {
+                displayName = ((DateTime)PropertyValue) != DateTime.MinValue ? Convert.ToString(PropertyValue, CultureInfo.CurrentCulture) : string.Empty;
+            }
+            else
+            {
+                displayName = Convert.ToString(PropertyValue, CultureInfo.CurrentCulture);
+            }
+            return displayName;
+        }
+    }
+}

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

@@ -5,33 +5,106 @@
 
 using System.Collections.Generic;
 using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts
 {
     /// <summary>
     /// Restore request parameters
     /// </summary>
-    public class RestoreParams
+    public class RestoreParams : GeneralRequestDetails
     {
+        public string SessionId
+        {
+            get
+            {
+                return GetOptionValue<string>("sessionId");
+            }
+            set
+            {
+                SetOptionValue("sessionId", value);
+            }
+        }
+
         /// <summary>
         /// The Uri to find the connection to do the restore operations
         /// </summary>
         public string OwnerUri { get; set; }
 
         /// <summary>
-        /// The backup file path
+        /// Comma delimited list of backup files
         /// </summary>
-        public string BackupFilePath { get; set; }
+        public string BackupFilePaths
+        {
+            get
+            {
+                return GetOptionValue<string>("backupFilePaths");
+            }
+            set
+            {
+                SetOptionValue("backupFilePaths", value);
+            }
+        }
 
         /// <summary>
         /// Target Database name to restore to
         /// </summary>
-        public string DatabaseName { get; set; }
+        public string TargetDatabaseName
+        {
+            get
+            {
+                return GetOptionValue<string>("targetDatabaseName");
+            }
+            set
+            {
+                SetOptionValue("targetDatabaseName", value);
+            }
+        }
+
+        /// <summary>
+        /// Source Database name to restore from
+        /// </summary>
+        public string SourceDatabaseName
+        {
+            get
+            {
+                return GetOptionValue<string>("sourceDatabaseName");
+            }
+            set
+            {
+                SetOptionValue("sourceDatabaseName", value);
+            }
+        }
 
         /// <summary>
         /// If set to true, the db files will be relocated to default data location in the server
         /// </summary>
-        public bool RelocateDbFiles { get; set; }
+        public bool RelocateDbFiles
+        {
+            get
+            {
+                return GetOptionValue<bool>("relocateDbFiles");
+            }
+            set
+            {
+                SetOptionValue("relocateDbFiles", value);
+            }
+        }
+
+        /// <summary>
+        /// Ids of the backup set to restore
+        /// </summary>
+        public string[] SelectedBackupSets
+        {
+            get
+            {
+                return GetOptionValue<string[]>("selectedBackupSets");
+            }
+            set
+            {
+                SetOptionValue("selectedBackupSets", value);
+            }
+        }
     }
 
     /// <summary>
@@ -87,10 +160,9 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts
     /// </summary>
     public class RestorePlanResponse
     {
-        /// <summary>
-        /// The backup file path
-        /// </summary>
-        public string BackupFilePath { get; set; }
+        public string RestoreSessionId { get; set; }
+
+        public DatabaseFileInfo[] BackupSetsToRestore { get; set; }
 
         /// <summary>
         /// Indicates whether the restore operation is supported 
@@ -107,6 +179,11 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts
         /// </summary>
         public IEnumerable<RestoreDatabaseFileInfo> DbFiles { get; set; }
 
+        /// <summary>
+        /// Database names extracted from backup sets
+        /// </summary>
+        public string[] DatabaseNamesFromBackupSets { get; set; }
+
         /// <summary>
         /// Server name
         /// </summary>

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

@@ -26,7 +26,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
     {
         private static readonly Lazy<DisasterRecoveryService> instance = new Lazy<DisasterRecoveryService>(() => new DisasterRecoveryService());
         private static ConnectionService connectionService = null;
-        private RestoreDatabaseHelper restoreDatabaseService = new RestoreDatabaseHelper();
+        private RestoreDatabaseHelper restoreDatabaseService = RestoreDatabaseHelper.Instance;
 
         /// <summary>
         /// Default, parameterless constructor.

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

@@ -3,6 +3,14 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
+using System.Globalization;
+using System.Linq;
+using Microsoft.SqlServer.Management.Smo;
+using Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts;
+
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 {
     /// <summary>
@@ -10,14 +18,186 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
     /// </summary>
     public class BackupSetInfo
     {
+        public const string BackupComponentPropertyName = "Component";
+        public const string NamePropertyName = "Name";
+        public const string BackupTypePropertyName = "Type";
+        public const string ServerNamePropertyName = "Server";
+        public const string DatabaseNamePropertyName = "Database";
+        public const string PositionPropertyName = "Position";
+        public const string FirstLsnPropertyName = "FirstLSN";
+        public const string LastLsnPropertyName = "LastLSN";
+        public const string CheckpointLsnPropertyName = "CheckpointLSN";
+        public const string FullLsnPropertyName = "FullLSN";
+        public const string StartDatePropertyName = "StartDate";
+        public const string FinishDatePropertyName = "FinishDate";
+        public const string SizePropertyName = "Size";
+        public const string UserNamePropertyName = "UserName";
+        public const string ExpirationPropertyName = "Expiration";
+
+        private Dictionary<string, LocalizedPropertyInfo> properties;
+
+        public BackupSetInfo(Dictionary<string, LocalizedPropertyInfo> properties)
+        {
+            this.properties = properties;
+        }
+
         /// <summary>
         /// Backup type (Full, Transaction Log, Differential ...)
         /// </summary>
-        public string BackupType { get; set; }
+        public string BackupType
+        {
+            get
+            {
+                return GetPropertyValueAsString(BackupTypePropertyName);
+            }
+        }
 
         /// <summary>
-        /// Backup component (Database, File, Log ...)
+        /// Backup set properties
         /// </summary>
-        public string BackupComponent { get; set; }
+        public ReadOnlyDictionary<string, LocalizedPropertyInfo> Properties
+        {
+            get
+            {
+                return new ReadOnlyDictionary<string, LocalizedPropertyInfo>(this.properties);
+            }
+        }
+
+        /// <summary>
+        /// Convert properties to array
+        /// </summary>
+        /// <returns></returns>
+        public LocalizedPropertyInfo[] ConvertPropertiesToArray()
+        {
+            return this.properties == null ? new LocalizedPropertyInfo[] { } : this.properties.Values.ToArray();
+        }
+
+        /// <summary>
+        /// Creates new BackupSet info
+        /// </summary>
+        /// <returns></returns>
+        public static BackupSetInfo Create(Restore restore, Server server)
+        {
+            BackupSet backupSet = restore.BackupSet;
+            Dictionary<string, LocalizedPropertyInfo> properties = new Dictionary<string, LocalizedPropertyInfo>();
+
+            string bkSetComponent;
+            string bkSetType;
+            CommonUtilities.GetBackupSetTypeAndComponent(backupSet.BackupSetType, out bkSetType, out bkSetComponent);
+
+            if (server.Version.Major > 8 && backupSet.IsCopyOnly)
+            {
+                bkSetType += SR.RestoreCopyOnly;
+            }
+
+            properties.Add(NamePropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = NamePropertyName,
+                PropertyValue = backupSet.Name,
+                PropertyDisplayName = SR.RestoreBackupSetName
+            });
+            properties.Add(BackupComponentPropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = BackupComponentPropertyName,
+                PropertyValue = bkSetComponent,
+                PropertyDisplayName = SR.RestoreBackupSetType
+            });
+            properties.Add(BackupTypePropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = BackupTypePropertyName,
+                PropertyValue = bkSetType,
+                PropertyDisplayName = SR.RestoreBackupSetComponent
+            });
+            properties.Add(ServerNamePropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = ServerNamePropertyName,
+                PropertyValue = backupSet.ServerName,
+                PropertyDisplayName = SR.RestoreBackupSetServer
+            });
+            properties.Add(DatabaseNamePropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = DatabaseNamePropertyName,
+                PropertyValue = backupSet.DatabaseName,
+                PropertyDisplayName = SR.RestoreBackupSetDatabase
+            });
+            properties.Add(PositionPropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = PositionPropertyName,
+                PropertyValueDisplayName = Convert.ToString(backupSet.Position, CultureInfo.CurrentCulture),
+                PropertyValue = backupSet.Position,
+                PropertyDisplayName = SR.RestoreBackupSetPosition
+            });
+            properties.Add(FirstLsnPropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = FirstLsnPropertyName,
+                PropertyValue = backupSet.FirstLsn,
+                PropertyDisplayName = SR.RestoreBackupSetFirstLsn
+            });
+            properties.Add(LastLsnPropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = LastLsnPropertyName,
+                PropertyValue = backupSet.LastLsn,
+                PropertyDisplayName = SR.RestoreBackupSetLastLsn
+            });
+            properties.Add(FullLsnPropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = FullLsnPropertyName,
+                PropertyValue = backupSet.DatabaseBackupLsn,
+                PropertyDisplayName = SR.RestoreBackupSetFullLsn
+            });
+            properties.Add(CheckpointLsnPropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = CheckpointLsnPropertyName,
+                PropertyValue = backupSet.CheckpointLsn,
+                PropertyDisplayName = SR.RestoreBackupSetCheckpointLsn
+            });
+            properties.Add(StartDatePropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = StartDatePropertyName,
+                PropertyValue = backupSet.BackupStartDate,
+                PropertyDisplayName = SR.RestoreBackupSetStartDate
+            });
+            properties.Add(FinishDatePropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = FinishDatePropertyName,
+                PropertyValue = backupSet.BackupFinishDate,
+                PropertyDisplayName = SR.RestoreBackupSetFinishDate
+            });
+            properties.Add(SizePropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = SizePropertyName,
+                PropertyValue = backupSet.BackupSize,
+                PropertyDisplayName = SR.RestoreBackupSetSize,
+            });
+            properties.Add(UserNamePropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = UserNamePropertyName,
+                PropertyValue = backupSet.UserName,
+                PropertyDisplayName = SR.RestoreBackupSetUserName,
+            });
+            properties.Add(ExpirationPropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = ExpirationPropertyName,
+                PropertyValue = backupSet.ExpirationDate,
+                PropertyDisplayName = SR.RestoreBackupSetExpiration,
+            });
+            properties.Add(DatabaseFileInfo.IdPropertyName, new LocalizedPropertyInfo
+            {
+                PropertyName = DatabaseFileInfo.IdPropertyName,
+                PropertyValue = backupSet.BackupSetGuid
+            });
+
+            return new BackupSetInfo(properties);
+        }
+
+        public string GetPropertyValueAsString(string propertyName)
+        {
+            LocalizedPropertyInfo propertyValue = null;
+            if(!string.IsNullOrEmpty(propertyName) && Properties != null)
+            {
+                Properties.TryGetValue(propertyName, out propertyValue);
+            }
+            return propertyValue.PropertyValue != null ? propertyValue.PropertyValue.ToString() : string.Empty;
+        }
     }
 }

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

@@ -15,6 +15,7 @@ using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
 using Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts;
 using Microsoft.SqlTools.ServiceLayer.TaskServices;
 using Microsoft.SqlTools.Utility;
+using System.Collections.Concurrent;
 
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 {
@@ -24,6 +25,22 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
     public class RestoreDatabaseHelper
     {
 
+        private static RestoreDatabaseHelper instance = new RestoreDatabaseHelper();
+        private ConcurrentDictionary<string, RestoreDatabaseTaskDataObject> restoreSessions = new ConcurrentDictionary<string, RestoreDatabaseTaskDataObject>();
+
+        internal RestoreDatabaseHelper()
+        {
+
+        }
+
+        public static RestoreDatabaseHelper Instance
+        {
+            get
+            {
+                return instance;
+            }
+        }
+
         /// <summary>
         /// Create a backup task for execution and cancellation
         /// </summary>
@@ -45,7 +62,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                     {
                         if (restoreDataObject.IsValid)
                         {
-                            ExecuteRestore(restoreDataObject);
+                            ExecuteRestore(restoreDataObject, sqlTask);
                             result.TaskStatus = SqlTaskStatus.Succeeded;
                         }
                         else
@@ -135,7 +152,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         {
             RestorePlanResponse response = new RestorePlanResponse()
             {
-                DatabaseName = restoreDataObject.RestoreParams.DatabaseName
+                DatabaseName = restoreDataObject.RestoreParams.TargetDatabaseName
             };
             try
             {
@@ -145,6 +162,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 
                     if (restoreDataObject != null && restoreDataObject.IsValid)
                     {
+                        response.RestoreSessionId = restoreDataObject.SessionId;
                         response.DatabaseName = restoreDataObject.TargetDatabase;
                         response.DbFiles = restoreDataObject.DbFiles.Select(x => new RestoreDatabaseFileInfo
                         {
@@ -160,6 +178,9 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                             response.ErrorMessage = SR.RestoreNotSupported;
                         }
 
+                        response.BackupSetsToRestore = restoreDataObject.GetBackupSetInfo().Select(x => new DatabaseFileInfo(x.ConvertPropertiesToArray())).ToArray();
+                        var dbNames = restoreDataObject.GetSourceDbNames();
+                        response.DatabaseNamesFromBackupSets = dbNames == null ? new string[] { } : dbNames.ToArray();
                         response.RelocateFilesNeeded = !restoreDataObject.DbFilesLocationAreValid();
                         response.DefaultDataFolder = restoreDataObject.DefaultDataFileFolder;
                         response.DefaultLogFolder = restoreDataObject.DefaultLogFileFolder;
@@ -204,7 +225,9 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             if (restoreDataObject != null)
             {
                 var backupTypes = restoreDataObject.GetBackupSetInfo();
-                return backupTypes.Any(x => x.BackupType.StartsWith(RestoreConstants.TypeFull));
+                var selectedBackupSets = restoreDataObject.RestoreParams.SelectedBackupSets;
+                return backupTypes.Any(x => (selectedBackupSets == null || selectedBackupSets.Contains(x.GetPropertyValueAsString(DatabaseFileInfo.IdPropertyName))) 
+                && x.BackupType.StartsWith(RestoreConstants.TypeFull));
             }
             return false;
         }
@@ -215,6 +238,24 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// <param name="restoreParams">Restore request parameters</param>
         /// <returns>Restore task object</returns>
         public RestoreDatabaseTaskDataObject CreateRestoreDatabaseTaskDataObject(RestoreParams restoreParams)
+        {
+            RestoreDatabaseTaskDataObject restoreTaskObject = null;
+            if (!string.IsNullOrWhiteSpace(restoreParams.SessionId))
+            {
+                this.restoreSessions.TryGetValue(restoreParams.SessionId, out restoreTaskObject);
+            }
+
+            if (restoreTaskObject == null)
+            {
+                restoreTaskObject = CreateRestoreForNewSession(restoreParams);
+                string sessionId = string.IsNullOrWhiteSpace(restoreParams.SessionId) ? Guid.NewGuid().ToString() : restoreParams.SessionId;
+                this.restoreSessions.AddOrUpdate(sessionId, restoreTaskObject, (key, oldSession) => restoreTaskObject);
+                restoreTaskObject.SessionId = sessionId;
+            }
+            return restoreTaskObject;
+        }
+
+        private RestoreDatabaseTaskDataObject CreateRestoreForNewSession(RestoreParams restoreParams)
         {
             ConnectionInfo connInfo;
             DisasterRecoveryService.ConnectionServiceInstance.TryFindConnection(
@@ -242,7 +283,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                 }
                 Server server = new Server(new ServerConnection(connection));
 
-                RestoreDatabaseTaskDataObject restoreDataObject = new RestoreDatabaseTaskDataObject(server, restoreParams.DatabaseName);
+                RestoreDatabaseTaskDataObject restoreDataObject = new RestoreDatabaseTaskDataObject(server, restoreParams.TargetDatabaseName);
                 restoreDataObject.RestoreParams = restoreParams;
                 return restoreDataObject;
             }
@@ -256,35 +297,54 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// <returns></returns>
         private void UpdateRestorePlan(RestoreDatabaseTaskDataObject restoreDataObject)
         {
-            if (!string.IsNullOrEmpty(restoreDataObject.RestoreParams.BackupFilePath))
+            if (restoreDataObject.PlanUpdateRequired)
             {
-                restoreDataObject.AddFile(restoreDataObject.RestoreParams.BackupFilePath);
-            }
-            restoreDataObject.RestorePlanner.ReadHeaderFromMedia = !string.IsNullOrEmpty(restoreDataObject.RestoreParams.BackupFilePath);
-           
-            restoreDataObject.RestorePlanner.DatabaseName = restoreDataObject.DefaultDbName;
-            restoreDataObject.TargetDatabase = restoreDataObject.RestoreParams.DatabaseName;
-            //TODO: used for other types of restore
-            /*bool isTailLogBackupPossible = restoreDataObject.RestorePlanner.IsTailLogBackupPossible(restoreDataObject.RestorePlanner.DatabaseName);
-            restoreDataObject.RestorePlanner.BackupTailLog = isTailLogBackupPossible;
-            restoreDataObject.TailLogBackupFile = restoreDataObject.Util.GetDefaultTailLogbackupFile(dbName);
-            restoreDataObject.RestorePlanner.TailLogBackupFile = restoreDataObject.TailLogBackupFile;
-            */
+                if (!string.IsNullOrEmpty(restoreDataObject.RestoreParams.BackupFilePaths))
+                {
+                    restoreDataObject.AddFiles(restoreDataObject.RestoreParams.BackupFilePaths);
+                }
+                restoreDataObject.RestorePlanner.ReadHeaderFromMedia = !string.IsNullOrEmpty(restoreDataObject.RestoreParams.BackupFilePaths);
 
-            restoreDataObject.UpdateRestorePlan(restoreDataObject.RestoreParams.RelocateDbFiles);
+                if (string.IsNullOrWhiteSpace(restoreDataObject.RestoreParams.SourceDatabaseName))
+                {
+                    restoreDataObject.RestorePlanner.DatabaseName = restoreDataObject.DefaultDbName;
+                }
+                else
+                {
+                    restoreDataObject.RestorePlanner.DatabaseName = restoreDataObject.RestoreParams.SourceDatabaseName;
+                }
+                restoreDataObject.TargetDatabase = restoreDataObject.RestoreParams.TargetDatabaseName;
+                //TODO: used for other types of restore
+                /*bool isTailLogBackupPossible = restoreDataObject.RestorePlanner.IsTailLogBackupPossible(restoreDataObject.RestorePlanner.DatabaseName);
+                restoreDataObject.RestorePlanner.BackupTailLog = isTailLogBackupPossible;
+                restoreDataObject.TailLogBackupFile = restoreDataObject.Util.GetDefaultTailLogbackupFile(dbName);
+                restoreDataObject.RestorePlanner.TailLogBackupFile = restoreDataObject.TailLogBackupFile;
+                */
+
+                restoreDataObject.UpdateRestorePlan(restoreDataObject.RestoreParams.RelocateDbFiles);
+            }
         }
 
         /// <summary>
         /// Executes the restore operation
         /// </summary>
         /// <param name="requestParam"></param>
-        public void ExecuteRestore(RestoreDatabaseTaskDataObject restoreDataObject)
+        public void ExecuteRestore(RestoreDatabaseTaskDataObject restoreDataObject, SqlTask sqlTask = null)
         {
+            // Restore Plan should be already created and updated at this point
             UpdateRestorePlan(restoreDataObject);
 
             if (restoreDataObject != null && CanRestore(restoreDataObject))
             {
-                restoreDataObject.RestorePlan.Execute();
+                try
+                {
+                    restoreDataObject.SqlTask = sqlTask;
+                    restoreDataObject.Execute();
+                }
+                catch(Exception ex)
+                {
+                    throw ex;
+                }
             }
             else
             {

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

@@ -9,6 +9,7 @@ using System.IO;
 using System.Linq;
 using Microsoft.SqlServer.Management.Smo;
 using Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts;
+using Microsoft.SqlTools.ServiceLayer.TaskServices;
 
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 {
@@ -17,8 +18,11 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
     /// </summary>
     public class RestoreDatabaseTaskDataObject
     {
+
+        private const char BackupMediaNameSeparator = ',';
         public RestoreDatabaseTaskDataObject(Server server, String databaseName)
         {
+            PlanUpdateRequired = true;
             this.Server = server;
             this.Util = new RestoreUtil(server);
             restorePlanner = new DatabaseRestorePlanner(server);
@@ -39,6 +43,16 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             this.restoreOptions.PercentCompleteNotification = 5;
         }
 
+        /// <summary>
+        /// Restore session id
+        /// </summary>
+        public string SessionId { get; set; }
+
+        /// <summary>
+        /// Sql task assigned to the restore object
+        /// </summary>
+        public SqlTask SqlTask { get; set; }
+
         public string TargetDatabase
         {
             get
@@ -86,14 +100,74 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// <summary>
         /// Add a backup file to restore plan media list
         /// </summary>
-        /// <param name="filePath"></param>
-        public void AddFile(string filePath)
+        /// <param name="filePaths"></param>
+        public void AddFiles(string filePaths)
         {
-            this.RestorePlanner.BackupMediaList.Add(new BackupDeviceItem
+            PlanUpdateRequired = true;
+            if (!string.IsNullOrWhiteSpace(filePaths))
             {
-                DeviceType = DeviceType.File,
-                Name = filePath
-            });
+                string[] files = filePaths.Split(BackupMediaNameSeparator);
+                files = files.Select(x => x.Trim()).ToArray();
+                foreach (var file in files)
+                {
+                    if (!this.RestorePlanner.BackupMediaList.Any(x => x.Name == file))
+                    {
+                        this.RestorePlanner.BackupMediaList.Add(new BackupDeviceItem
+                        {
+                            DeviceType = DeviceType.File,
+                            Name = file
+                        });
+                    }
+                }
+
+                var itemsToRemove = this.RestorePlanner.BackupMediaList.Where(x => !files.Contains(x.Name));
+                foreach (var item in itemsToRemove)
+                {
+                    this.RestorePlanner.BackupMediaList.Remove(item);
+                }
+            }
+        }
+
+        /// <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>
+        /// Executes the restore operations
+        /// </summary>
+        public void Execute()
+        {
+            RestorePlan restorePlan = RestorePlan;
+            // 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)
+                    {
+                        SqlTask.AddMessage($"{e.Percent}%", SqlTaskStatus.InProgress);
+                    }
+                };
+                restorePlan.Execute();
+            }
         }
 
         public RestoreUtil Util { get; set; }
@@ -109,7 +183,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         }
 
         private string tailLogBackupFile;
-        private bool planUpdateRequired = false;
+        public bool PlanUpdateRequired { get; private set; }
 
         /// <summary>
         /// File to backup tail log before doing the restore
@@ -122,7 +196,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                 if (tailLogBackupFile == null || !tailLogBackupFile.Equals(value))
                 {
                     this.RestorePlanner.TailLogBackupFile = value;
-                    this.planUpdateRequired = true;
+                    this.PlanUpdateRequired = true;
                     this.tailLogBackupFile = value;
                 }
             }
@@ -384,43 +458,6 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
         /// </summary>
         internal string ContainerSharedAccessPolicy = string.Empty;
 
-        /// <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.UpdateRestoreSelected();
-                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.RestoreSelected[i] == true)
-                    {
-                        rp.RestoreOperations.Add(res);
-                    }
-                    i++;
-                }
-            }
-            this.SetRestorePlanProperties(rp);
-            return rp;
-        }
-
         /// <summary>
         /// Updates the RestoreSelected Array to hold information about updated Restore Plan
         /// </summary>
@@ -472,17 +509,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
             List<BackupSetInfo> result = new List<BackupSetInfo>();
             foreach (Restore restore in RestorePlan.RestoreOperations)
             {
-                BackupSet backupSet = restore.BackupSet;
-
-                String bkSetComponent;
-                String bkSetType;
-                CommonUtilities.GetBackupSetTypeAndComponent(backupSet.BackupSetType, out bkSetType, out bkSetComponent);
-
-                if (this.Server.Version.Major > 8 && backupSet.IsCopyOnly)
-                {
-                    bkSetType += " (Copy Only)";
-                }
-                result.Add(new BackupSetInfo { BackupComponent = bkSetComponent, BackupType = bkSetType });
+                result.Add(BackupSetInfo.Create(restore, Server));
             }
 
             return result;
@@ -564,7 +591,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                                          //Clear any existing exceptions as new plan is getting recreated.
             this.CreateOrUpdateRestorePlanException = null;
             this.DbFiles.Clear();
-            this.planUpdateRequired = false;
+            this.PlanUpdateRequired = false;
             this.restorePlan = null;
             if (String.IsNullOrEmpty(this.RestorePlanner.DatabaseName))
             {
@@ -666,7 +693,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                 if (this.RestorePlanner.BackupTailLog != value)
                 {
                     this.RestorePlanner.BackupTailLog = value;
-                    this.planUpdateRequired = true;
+                    this.PlanUpdateRequired = true;
                 }
             }
         }
@@ -685,7 +712,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
                 if (this.RestorePlanner.TailLogWithNoRecovery != value)
                 {
                     this.RestorePlanner.TailLogWithNoRecovery = value;
-                    this.planUpdateRequired = true;
+                    this.PlanUpdateRequired = true;
                 }
             }
         }