From f95d8cb40272d08936263567aa94479d7cf58e52 Mon Sep 17 00:00:00 2001
From: srayan <sguhat@microsoft.com>
Date: Wed, 21 Feb 2024 10:12:26 -0800
Subject: [PATCH] GetStorageAccountAsync is making a Storage Account management
 API call, which is restricted at 800 requests / 5 minutes. This is different
 from calls made to the storage account data place (which have a much higher
 threshold).
 https://learn.microsoft.com/en-us/azure/role-based-access-control/resource-provider-operations#microsoftstorage
 has list of APIs.

---
 migrationTool/ams/AssetMigrator.cs | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/migrationTool/ams/AssetMigrator.cs b/migrationTool/ams/AssetMigrator.cs
index 32eab08..b21c5e8 100644
--- a/migrationTool/ams/AssetMigrator.cs
+++ b/migrationTool/ams/AssetMigrator.cs
@@ -6,6 +6,7 @@
 using Azure.Storage.Blobs;
 using Microsoft.Extensions.Logging;
 using Spectre.Console;
+using System.Collections.Concurrent;
 using System.Diagnostics;
 using System.Threading.Channels;
 
@@ -16,6 +17,7 @@ internal class AssetMigrator : BaseMigrator
         private readonly TransformFactory _transformFactory;
         private readonly AssetOptions _options;
         private readonly IMigrationTracker<BlobContainerClient, AssetMigrationResult> _tracker;
+        protected readonly TokenCredential _credentials;
 
         public AssetMigrator(
             GlobalOptions globalOptions,
@@ -30,6 +32,7 @@ public AssetMigrator(
             _options = assetOptions;
             _tracker = tracker;
             _transformFactory = transformFactory;
+            _credentials = credential;
         }
 
         public override async Task MigrateAsync(CancellationToken cancellationToken)
@@ -83,11 +86,22 @@ public override async Task MigrateAsync(CancellationToken cancellationToken)
         private async Task<AssetStats> MigrateAsync(MediaServicesAccountResource account, AsyncPageable<MediaAssetResource> assets, List<MediaAssetResource>? filteredList, ChannelWriter<double> writer, CancellationToken cancellationToken)
         {
             var stats = new AssetStats();
+
+            // Dictionary to restrict spawning of too many BlobServiceClients
+            var storageAccounts = new ConcurrentDictionary<string, BlobServiceClient>();
+
             await MigrateInParallel(assets, filteredList, async (asset, cancellationToken) =>
             {
-                var storage = await _resourceProvider.GetStorageAccountAsync(account, asset, cancellationToken);
+                var storageAccountName = asset.Data.StorageAccountName;
+                if (!storageAccounts.ContainsKey(storageAccountName))
+                {
+                    BlobServiceClient newBlobServiceClient = new BlobServiceClient(new Uri($"https://{storageAccountName}.blob.core.windows.net"), _credentials);
+                    storageAccounts.TryAdd(storageAccountName, newBlobServiceClient);
+                }
+
+                BlobServiceClient blobServiceClient = storageAccounts[storageAccountName];
 
-                var result = await MigrateAsync(account, storage, asset, cancellationToken);
+                var result = await MigrateAsync(account, blobServiceClient, asset, cancellationToken);
                 stats.Update(result);
                 await writer.WriteAsync(stats.Total, cancellationToken);