From 116119f91a158bc7ddf2c17fd2ddf68824fd37b6 Mon Sep 17 00:00:00 2001
From: Erick Yondon <8766776+erdembayar@users.noreply.github.com>
Date: Mon, 26 Aug 2024 20:15:29 -0700
Subject: [PATCH] Add Table storage MSI

---
 .../NuGet.Jobs.Common.csproj                  |  1 +
 .../StorageAccountExtensions.cs               | 69 +++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/src/NuGet.Jobs.Common/NuGet.Jobs.Common.csproj b/src/NuGet.Jobs.Common/NuGet.Jobs.Common.csproj
index 99fc47c49b..2acb102c93 100644
--- a/src/NuGet.Jobs.Common/NuGet.Jobs.Common.csproj
+++ b/src/NuGet.Jobs.Common/NuGet.Jobs.Common.csproj
@@ -13,6 +13,7 @@
 
   <ItemGroup>
     <PackageReference Include="Autofac.Extensions.DependencyInjection" />
+    <PackageReference Include="Azure.Data.Tables" />
     <PackageReference Include="Dapper.StrongName" />
     <PackageReference Include="Microsoft.Extensions.DependencyInjection" />
     <PackageReference Include="Microsoft.Extensions.Options.ConfigurationExtensions" />
diff --git a/src/NuGet.Jobs.Common/StorageAccountExtensions.cs b/src/NuGet.Jobs.Common/StorageAccountExtensions.cs
index 1671145bbd..0e2d0d39bf 100644
--- a/src/NuGet.Jobs.Common/StorageAccountExtensions.cs
+++ b/src/NuGet.Jobs.Common/StorageAccountExtensions.cs
@@ -4,6 +4,8 @@
 using System;
 using Autofac;
 using Autofac.Builder;
+using Azure.Data.Tables;
+using Azure.Identity;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
@@ -103,6 +105,50 @@ public static IRegistrationBuilder<CloudBlobClientWrapper, SimpleActivatorData,
             });
         }
 
+        public static TableServiceClient CreateTableServiceClient(
+            this IServiceProvider serviceProvider,
+            string storageConnectionString)
+        {
+            if (serviceProvider == null)
+            {
+                throw new ArgumentNullException(nameof(serviceProvider));
+            }
+            if (string.IsNullOrWhiteSpace(storageConnectionString))
+            {
+                throw new ArgumentException($"{nameof(storageConnectionString)} cannot be null or empty.", nameof(storageConnectionString));
+            }
+
+            var msiConfiguration = serviceProvider.GetRequiredService<IOptions<StorageMsiConfiguration>>().Value;
+            return CreateTableServiceClientClient(
+                msiConfiguration,
+                storageConnectionString);
+        }
+
+        public static IRegistrationBuilder<TableServiceClient, SimpleActivatorData, SingleRegistrationStyle> RegisterTableServiceClient<TConfiguration>(
+            this ContainerBuilder builder,
+            Func<TConfiguration, string> getConnectionString)
+            where TConfiguration : class, new()
+        {
+            if (builder == null)
+            {
+                throw new ArgumentNullException(nameof(builder));
+            }
+            if (getConnectionString == null)
+            {
+                throw new ArgumentNullException(nameof(getConnectionString));
+            }
+
+            return builder.Register(c =>
+            {
+                var options = c.Resolve<IOptionsSnapshot<TConfiguration>>();
+                string storageConnectionString = getConnectionString(options.Value);
+                var msiConfiguration = c.Resolve<IOptions<StorageMsiConfiguration>>().Value;
+                return CreateTableServiceClientClient(
+                    msiConfiguration,
+                    storageConnectionString);
+            });
+        }
+
         private static CloudBlobClientWrapper CreateCloudBlobClient(
             StorageMsiConfiguration msiConfiguration,
             string storageConnectionString,
@@ -133,5 +179,28 @@ private static CloudBlobClientWrapper CreateCloudBlobClient(
                 readAccessGeoRedundant,
                 requestTimeout);
         }
+
+        private static TableServiceClient CreateTableServiceClientClient(
+            StorageMsiConfiguration msiConfiguration,
+            string tableStorageConnectionString)
+        {
+            if (msiConfiguration.UseManagedIdentity)
+            {
+                if (string.IsNullOrWhiteSpace(msiConfiguration.ManagedIdentityClientId))
+                {
+                    return new TableServiceClient(new Uri(tableStorageConnectionString),
+                        new DefaultAzureCredential());
+                }
+                else
+                {
+                    return new TableServiceClient(new Uri(tableStorageConnectionString),
+                        new ManagedIdentityCredential(msiConfiguration.ManagedIdentityClientId));
+                }
+            }
+
+            // workaround for https://github.com/Azure/azure-sdk-for-net/issues/44373
+            tableStorageConnectionString.Replace("SharedAccessSignature=?", "SharedAccessSignature=");
+            return new TableServiceClient(tableStorageConnectionString);
+        }
     }
 }