Add support for cross tenant and cross subscription restore operations (#17582)

* Add logic to support cross server and cross tenant restore

1) Added support for cross restore operations by using optional param sourceResourceId.
2) Added support for cross tenant restore operations.
3) Reference the new 2.1.1-preview .NET sdk, as it has the newly added sourceResourceId param.

* Add release notes to changelog.md

* Add null check for parser, modify release notes

Author: imadan1996

src/Sql/Sql.Test/Sql.Test.csproj

@@ -19,7 +19,7 @@
     <PackageReference Include="Microsoft.Azure.KeyVault.WebKey" Version="3.0.1" />
     <PackageReference Include="Microsoft.Azure.Management.KeyVault" Version="4.0.0-preview.1" />
     <PackageReference Include="Microsoft.Azure.Management.OperationalInsights" Version="0.24.0-preview" />
-    <PackageReference Include="Microsoft.Azure.Management.Sql" Version="2.1.0-preview" />
+    <PackageReference Include="Microsoft.Azure.Management.Sql" Version="2.1.1-preview" />
   </ItemGroup>
 
   <ItemGroup>

src/Sql/Sql/ChangeLog.md

@@ -31,6 +31,7 @@
     - `Get-AzSqlInstanceLink`
     - `Remove-AzSqlInstanceLink`
     - `Set-AzSqlInstanceLink`
+* Added support for DataWarehouse cross tenant and cross subscription restore operations to `Restore-AzSqlDatabase` cmdlet
 
 ## Version 3.7.1
 * Deprecation of Get-AzSqlDatabaseTransparentDataEncryptionActivity cmdlet

src/Sql/Sql/Database Backup/Cmdlet/RestoreAzureRMSqlDatabase.cs

@@ -20,6 +20,7 @@
 using Microsoft.WindowsAzure.Commands.Utilities.Common;
 using System;
 using System.Collections;
+using System.Collections.Generic;
 using System.Globalization;
 using System.Linq;
 using System.Management.Automation;
@@ -398,7 +399,45 @@ protected override AzureSqlDatabaseModel GetEntity()
                 model.Edition = Edition;
             }
 
-            return ModelAdapter.RestoreDatabase(this.ResourceGroupName, restorePointInTime, ResourceId, model);
+            /// get auth headers for cross-sub and cross-tenant restore operations
+            string targetSubscriptionId = ModelAdapter.Context?.Subscription.Id;
+            string sourceSubscriptionId = ParseSourceSubscriptionIdFromResourceId(ResourceId);
+            Dictionary<string, List<string>> auxAuthHeader = null;
+            if (!string.IsNullOrEmpty(ResourceId) && targetSubscriptionId!=sourceSubscriptionId)
+            {
+                List<string> resourceIds = new List<string>();
+                resourceIds.Add(ResourceId);
+                var auxHeaderDictionary = GetAuxilaryAuthHeaderFromResourceIds(resourceIds);
+                if (auxHeaderDictionary != null && auxHeaderDictionary.Count > 0)
+                {
+                    auxAuthHeader = new Dictionary<string, List<string>>(auxHeaderDictionary);
+                }
+            }
+
+            return ModelAdapter.RestoreDatabase(this.ResourceGroupName, restorePointInTime, ResourceId, model, sourceSubscriptionId, auxAuthHeader);
+        }
+
+        /// <summary>
+        /// Parse source subscription id from ResourceId
+        /// </summary>
+        /// <returns>Source Subscription Id</returns>
+        private string ParseSourceSubscriptionIdFromResourceId(string resourceId)
+        {
+            if (string.IsNullOrEmpty(resourceId))
+            {
+            return null;
+            }
+
+            string[] words = resourceId.Split('/');
+            string sourceSubscriptionId = "";
+            for (int i = 0; i < words.Length; i++)
+            {
+                if (words[i] == "subscriptions")
+                {
+                    sourceSubscriptionId = words[i + 1];
+                }
+            }
+            return sourceSubscriptionId;
         }
     }
 }

src/Sql/Sql/Database Backup/Services/AzureSqlDatabaseBackupAdapter.cs

@@ -13,17 +13,12 @@
 // ----------------------------------------------------------------------------------
 
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
-using Microsoft.Azure.Commands.Common.Authentication.Models;
 using Microsoft.Azure.Commands.Sql.Backup.Model;
 using Microsoft.Azure.Commands.Sql.Database.Model;
-using Microsoft.Azure.Commands.Sql.Database.Services;
-using Microsoft.Azure.Commands.Sql.ElasticPool.Services;
 using Microsoft.Azure.Commands.Sql.Server.Adapter;
 using Microsoft.Azure.Management.Sql.LegacySdk.Models;
-using Microsoft.WindowsAzure.Commands.Utilities.Common;
 using System;
 using System.Collections.Generic;
-using System.Globalization;
 using System.Linq;
 using System.Management.Automation;
 
@@ -586,8 +581,10 @@ internal AzureSqlDatabaseGeoBackupPolicyModel SetDatabaseGeoBackupPolicy(
         /// <param name="restorePointInTime">A point to time to restore to (for PITR and dropped DB restore)</param>
         /// <param name="resourceId">The resource ID of the DB to restore (live, geo backup, deleted database, long term retention backup, etc.)</param>
         /// <param name="model">An object modeling the database to create via restore</param>
+        /// <param name="sourceSubscriptionId">Source Subscription Id</param>
+        /// <param name="customHeaders">Custom headers</param>
         /// <returns>Restored database object</returns>
-        internal AzureSqlDatabaseModel RestoreDatabase(string resourceGroup, DateTime restorePointInTime, string resourceId, AzureSqlDatabaseModel model)
+        internal AzureSqlDatabaseModel RestoreDatabase(string resourceGroup, DateTime restorePointInTime, string resourceId, AzureSqlDatabaseModel model, string sourceSubscriptionId, Dictionary<string, List<string>> customHeaders = null)
         {
             // Construct the ARM resource Id of the pool
             string elasticPoolId = string.IsNullOrWhiteSpace(model.ElasticPoolName) ? null : AzureSqlDatabaseModel.PoolIdTemplate.FormatInvariant(
@@ -614,29 +611,46 @@ internal AzureSqlDatabaseModel RestoreDatabase(string resourceGroup, DateTime re
                 RequestedBackupStorageRedundancy = model.RequestedBackupStorageRedundancy,
                 ZoneRedundant = model.ZoneRedundant,
             };
-
-            if (model.CreateMode == Management.Sql.Models.CreateMode.Recovery)
-            {
-                dbModel.RecoverableDatabaseId = resourceId;
-            }
-            else if (model.CreateMode == Management.Sql.Models.CreateMode.Restore)
-            {
-                dbModel.RestorableDroppedDatabaseId = resourceId;
-            }
-            else if (model.CreateMode == Management.Sql.Models.CreateMode.PointInTimeRestore)
+            
+            // check if restore operation is cross subscription or same subscription
+            if (_subscription.Id != sourceSubscriptionId)
             {
-                dbModel.SourceDatabaseId = resourceId;
-            }
-            else if (model.CreateMode == Management.Sql.Models.CreateMode.RestoreLongTermRetentionBackup)
-            {
-                dbModel.LongTermRetentionBackupResourceId = resourceId;
+                // cross subscription path
+                if (dbModel.CreateMode != Management.Sql.Models.CreateMode.Recovery 
+                    && dbModel.CreateMode != Management.Sql.Models.CreateMode.Restore 
+                    && dbModel.CreateMode != Management.Sql.Models.CreateMode.PointInTimeRestore 
+                    && dbModel.CreateMode != Management.Sql.Models.CreateMode.RestoreLongTermRetentionBackup)
+				{
+                    throw new ArgumentException("Restore mode not supported");
+                }
+                dbModel.SourceResourceId = resourceId;
             }
             else
             {
-                throw new ArgumentException("Restore mode not supported");
+                // same subscription path
+                if (model.CreateMode == Management.Sql.Models.CreateMode.Recovery)
+                {
+                    dbModel.RecoverableDatabaseId = resourceId;
+                }
+                else if (model.CreateMode == Management.Sql.Models.CreateMode.Restore)
+                {
+                    dbModel.RestorableDroppedDatabaseId = resourceId;
+                }
+                else if (model.CreateMode == Management.Sql.Models.CreateMode.PointInTimeRestore)
+                {
+                    dbModel.SourceDatabaseId = resourceId;
+                }
+                else if (model.CreateMode == Management.Sql.Models.CreateMode.RestoreLongTermRetentionBackup)
+                {
+                    dbModel.LongTermRetentionBackupResourceId = resourceId;
+                }
+                else
+                {
+                    throw new ArgumentException("Restore mode not supported");
+                }
             }
 
-            Management.Sql.Models.Database database = Communicator.RestoreDatabase(resourceGroup, model.ServerName, model.DatabaseName, dbModel);
+            Management.Sql.Models.Database database = Communicator.RestoreDatabase(resourceGroup, model.ServerName, model.DatabaseName, dbModel, customHeaders);
 
             return new AzureSqlDatabaseModel(resourceGroup, model.ServerName, database);
         }

src/Sql/Sql/Database Backup/Services/AzureSqlDatabaseBackupCommunicator.cs

@@ -19,8 +19,6 @@
 using Microsoft.Azure.Management.Sql.LegacySdk;
 using Microsoft.Azure.Management.Sql.LegacySdk.Models;
 using System.Collections.Generic;
-using Microsoft.Azure.Commands.Sql.Database.Model;
-using Microsoft.Azure.Management.Internal.Resources.Models;
 
 namespace Microsoft.Azure.Commands.Sql.Backup.Services
 {
@@ -461,10 +459,23 @@ public Management.Sql.LegacySdk.Models.Database LegacyRestoreDatabase(string res
         /// <param name="serverName">The name of the Azure SQL Server</param>
         /// <param name="databaseName">The name of the Azure SQL database</param>
         /// <param name="model">Sql Database Model with required parameters</param>
+        /// <param name="customHeaders">Custom headers</param>
         /// <returns>Restored database object</returns>
-        public Management.Sql.Models.Database RestoreDatabase(string resourceGroupName, string serverName, string databaseName, Management.Sql.Models.Database model)
+        public Management.Sql.Models.Database RestoreDatabase(string resourceGroupName, string serverName, string databaseName, Management.Sql.Models.Database model, Dictionary<string, List<string>> customHeaders = null)
         {
-            return GetCurrentSqlClient().Databases.CreateOrUpdate(resourceGroupName, serverName, databaseName, model);
+            if (customHeaders == null)
+            {
+                // Execute the create call without the custom headers. 
+                return GetCurrentSqlClient().Databases.CreateOrUpdate(resourceGroupName, serverName, databaseName, model);
+            }
+            else
+            {
+                // Execute the create call and pass the custom headers. 
+                using (var _result = GetCurrentSqlClient().Databases.CreateOrUpdateWithHttpMessagesAsync(resourceGroupName, serverName, databaseName, model, customHeaders).ConfigureAwait(false).GetAwaiter().GetResult())
+                {
+                    return _result.Body;
+                }
+            }
         }
 
         /// <summary>

src/Sql/Sql/Sql.csproj

@@ -21,7 +21,7 @@
   </ItemGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Azure.Management.Sql" Version="2.1.0-preview" />
+    <PackageReference Include="Microsoft.Azure.Management.Sql" Version="2.1.1-preview" />
     <PackageReference Include="System.Security.Permissions" Version="4.5.0" />
   </ItemGroup>
   <ItemGroup>