Enhancement: handle multiple resource groups in one AMS account (#172)

* Enhancement:handle multiple resource groups in one AMS account * update * address comments * update * updates * update
Azure · Sep 8, 2023 · 1a70d8c · 1a70d8c
1 parent 4f70f9a
commit 1a70d8c
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 12 deletions.
diff --git a/ams/AssetAnalyzer.cs b/ams/AssetAnalyzer.cs
@@ -115,6 +115,7 @@ public override async Task MigrateAsync(CancellationToken cancellationToken)
                 reportGenerator = new ReportGenerator(file);
                 reportGenerator.WriteHeader();
             }
+            await _resourceProvider.SetStorageResourceGroupsAsync(account, cancellationToken);
             var assets = account.GetMediaAssets()
                 .GetAllAsync(resourceFilter, cancellationToken: cancellationToken);
             var statistics = new AssetStats();
@@ -139,7 +140,6 @@ public override async Task MigrateAsync(CancellationToken cancellationToken)
             await MigrateInParallel(assets, filteredList, async (asset, cancellationToken) =>
             {
                 var storage = await _resourceProvider.GetStorageAccountAsync(account, asset, cancellationToken);
-
                 var result = await AnalyzeAsync(asset, storage, cancellationToken);
                 var assetType = result.AssetType ?? "unknown";
                 assetTypes.AddOrUpdate(assetType, 1, (key, value) => Interlocked.Increment(ref value));

diff --git a/ams/AssetMigrator.cs b/ams/AssetMigrator.cs
@@ -51,8 +51,9 @@ public override async Task MigrateAsync(CancellationToken cancellationToken)
                                                         _options.CreationTimeEnd);
 
             var orderBy = "properties/created";
-            var assets = account.GetMediaAssets()
-                .GetAllAsync(resourceFilter, orderby: orderBy, cancellationToken: cancellationToken);
+
+            await _resourceProvider.SetStorageResourceGroupsAsync(account, cancellationToken);
+            var assets = account.GetMediaAssets().GetAllAsync(resourceFilter, orderby: orderBy, cancellationToken: cancellationToken);
 
             List<MediaAssetResource>? filteredList = null;
 
@@ -82,9 +83,11 @@ private async Task<AssetStats> MigrateAsync(MediaServicesAccountResource account
             await MigrateInParallel(assets, filteredList, async (asset, cancellationToken) =>
             {
                 var storage = await _resourceProvider.GetStorageAccountAsync(account, asset, cancellationToken);
+
                 var result = await MigrateAsync(account, storage, asset, cancellationToken);
                 stats.Update(result);
                 await writer.WriteAsync(stats.Total, cancellationToken);
+
             },
             _options.BatchSize,
             cancellationToken);

diff --git a/ams/AzureResourceProvider.cs b/ams/AzureResourceProvider.cs
@@ -2,9 +2,10 @@
 using AMSMigrate.Contracts;
 using Azure.Core;
 using Azure.ResourceManager;
-using Azure.ResourceManager.Storage;
 using Azure.ResourceManager.Media;
+using Azure.ResourceManager.Media.Models;
 using Azure.ResourceManager.Resources;
+using Azure.ResourceManager.Storage;
 using Azure.Storage.Blobs;
 
 namespace AMSMigrate.Ams
@@ -14,35 +15,72 @@ class AzureResourceProvider
         protected readonly ResourceGroupResource _resourceGroup;
         protected readonly GlobalOptions _globalOptions;
         protected readonly TokenCredential _credentials;
+        protected readonly ArmClient _armClient;
+
+        private Dictionary<string, ResourceGroupResource> _storageResourceGroups;
 
         public AzureResourceProvider(TokenCredential credential, GlobalOptions options)
         {
             _globalOptions = options;
             _credentials = credential;
             var clientOptions = new ArmClientOptions();
             clientOptions.Diagnostics.ApplicationId = $"AMSMigrate/{GetType().Assembly.GetName().Version}";
-            var armClient = new ArmClient(credential, default, clientOptions);
+            _armClient = new ArmClient(credential, default, clientOptions);
             var resourceGroupId = ResourceGroupResource.CreateResourceIdentifier(
                 options.SubscriptionId,
                 options.ResourceGroup);
-            _resourceGroup = armClient.GetResourceGroupResource(resourceGroupId);
+            _resourceGroup = _armClient.GetResourceGroupResource(resourceGroupId);
+            _storageResourceGroups = new Dictionary<string, ResourceGroupResource>();
+        }
+
+        public async Task SetStorageResourceGroupsAsync(MediaServicesAccountResource account, CancellationToken cancellationToken)
+        {
+            IList<MediaServicesStorageAccount> storageAccounts;
+            var mediaServiceResource = await account.GetAsync(cancellationToken: cancellationToken);
+            if (mediaServiceResource.GetRawResponse().Status == 200)
+            {
+                storageAccounts = mediaServiceResource.Value.Data.StorageAccounts;
+                if (storageAccounts != null && storageAccounts.Any())
+                {
+                    foreach (var storageAccount in storageAccounts)
+                    {
+                        string? storageAccountId = storageAccount.Id;  
+                        ///subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/microsoft.storage/storageaccounts/{accountName}
+                        if (!string.IsNullOrEmpty(storageAccountId))
+                        {
+                            string[] parts;
+                            parts = storageAccountId.Split('/');
+                            string resourceGroupName = parts[4];
+                            string storageAccName = parts[8];
+                            string subscriptionId = parts[2];
+                            var resourceGroupId = ResourceGroupResource.CreateResourceIdentifier(
+                                 subscriptionId,
+                                 resourceGroupName);
+                            var resourceGroup = _armClient.GetResourceGroupResource(resourceGroupId);
+                            _storageResourceGroups.Add(storageAccName, resourceGroup);
+                       }
+                    }
+                }
+
+            }
         }
 
         public async Task<MediaServicesAccountResource> GetMediaAccountAsync(
             string mediaAccountName,
             CancellationToken cancellationToken)
         {
             return await _resourceGroup.GetMediaServicesAccountAsync(
-                mediaAccountName, cancellationToken);
-        }
+               mediaAccountName, cancellationToken);
+       }
 
         public async Task<BlobServiceClient> GetStorageAccountAsync(
             MediaServicesAccountResource account,
             MediaAssetResource asset,
             CancellationToken cancellationToken)
         {
             string assetStorageAccountName = asset.Data.StorageAccountName;
-            var resource = await _resourceGroup.GetStorageAccountAsync(asset.Data.StorageAccountName, cancellationToken: cancellationToken);
+            _storageResourceGroups.TryGetValue(asset.Data.StorageAccountName, out var rg);
+            var resource = await rg.GetStorageAccountAsync(asset.Data.StorageAccountName, cancellationToken: cancellationToken);
             return GetStorageAccount(resource);
         }
 

diff --git a/ams/CleanupCommand.cs b/ams/CleanupCommand.cs
@@ -52,6 +52,7 @@ public override async Task MigrateAsync(CancellationToken cancellationToken)
             var resourceFilter = _options.IsCleanUpAccount ? null : GetAssetResourceFilter(_options.ResourceFilter, null, null);
 
             var orderBy = "properties/created";
+            await _resourceProvider.SetStorageResourceGroupsAsync(account, cancellationToken);
             assets = account.GetMediaAssets()
                 .GetAllAsync(resourceFilter, orderby: orderBy, cancellationToken: cancellationToken);
             List<MediaAssetResource>? assetList = await assets.ToListAsync(cancellationToken);
@@ -161,10 +162,8 @@ private async Task<bool> CleanUpAssetAsync(bool isForcedelete, MediaServicesAcco
                 if (!await container.ExistsAsync(cancellationToken))
                 {
                     _logger.LogWarning("Container {name} missing for asset {asset}", container.Name, asset.Data.Name);
-
                     return false;
                 }
-
                 // The asset container exists, try to check the metadata list first.
 
                 if (isForcedelete || (_tracker.GetMigrationStatusAsync(container, cancellationToken).Result.Status == MigrationStatus.Completed))

diff --git a/ams/ResetCommand.cs b/ams/ResetCommand.cs
@@ -36,7 +36,7 @@ public override async Task MigrateAsync(CancellationToken cancellationToken)
         {
             var account = await GetMediaAccountAsync(_options.AccountName, cancellationToken);
             _logger.LogInformation("Begin reset assets on account: {name}", account.Data.Name);
-
+            await _resourceProvider.SetStorageResourceGroupsAsync(account, cancellationToken);
             AsyncPageable<MediaAssetResource> assets = account.GetMediaAssets()
                 .GetAllAsync(cancellationToken: cancellationToken);
             List<MediaAssetResource>? assetList = await assets.ToListAsync(cancellationToken);