Enable SQL Workload Identity (#3085)

- Allow use of explicit workload identity for SQL connections
- Enable configuring retries for Key Vault credentials

src/Microsoft.Health.Dicom.Azure/Registration/KeyVaultClientRegistrationExtensions.cs

@@ -10,6 +10,7 @@
 using Microsoft.Extensions.Azure;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Health.Core.Extensions;
 using Microsoft.Health.Dicom.Azure;
 using Microsoft.Health.Dicom.Azure.KeyVault;
 using Microsoft.Health.Dicom.Core.Features.Common;
@@ -91,7 +92,9 @@ private static IServiceCollection AddKeyVaultClient(
 
             services.AddAzureClients(builder =>
             {
-                IAzureClientBuilder<SecretClient, SecretClientOptions> clientBuilder = builder.AddSecretClient(section);
+                IAzureClientBuilder<SecretClient, SecretClientOptions> clientBuilder = builder
+                    .AddSecretClient(section)
+                    .WithRetryableCredential(section);
 
                 if (configureOptions != null)
                     clientBuilder.ConfigureOptions(configureOptions);

src/Microsoft.Health.Dicom.SqlServer.UnitTests/Microsoft.Health.Dicom.SqlServer.UnitTests.csproj

@@ -7,11 +7,15 @@
   <ItemGroup>
     <PackageReference Include="fo-dicom" />
     <PackageReference Include="Microsoft.Data.SqlClient" />
+    <PackageReference Include="Microsoft.Extensions.Configuration" />
+    <PackageReference Include="Microsoft.Extensions.Configuration.Abstractions" />
+    <PackageReference Include="Microsoft.Extensions.DependencyInjection" />
+    <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" />
     <PackageReference Include="Microsoft.Health.SqlServer" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" />
+    <PackageReference Include="System.Drawing.Common" PrivateAssets="All" />
     <PackageReference Include="xunit" />
     <PackageReference Include="xunit.runner.visualstudio" PrivateAssets="All" />
-    <PackageReference Include="System.Drawing.Common" PrivateAssets="All" />
   </ItemGroup>
 
   <ItemGroup>

src/Microsoft.Health.Dicom.SqlServer.UnitTests/Registration/DicomSqlServerRegistrationExtensionsTests.cs

@@ -0,0 +1,98 @@
+// -------------------------------------------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+// -------------------------------------------------------------------------------------------------
+
+using System;
+using System.Collections.Generic;
+using Microsoft.Data.SqlClient;
+using Microsoft.Extensions.Configuration;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Health.Dicom.Core.Registration;
+using Microsoft.Health.Dicom.SqlServer.Registration;
+using Microsoft.Health.SqlServer.Configs;
+using NSubstitute;
+using Xunit;
+
+namespace Microsoft.Health.Dicom.SqlServer.UnitTests.Registration;
+
+public sealed class DicomSqlServerRegistrationExtensionsTests : IDisposable
+{
+    [Fact]
+    public void GivenWebServerBuilder_WhenUsingDefaults_ThenUseBuiltInAuthenticationProvider()
+    {
+        IServiceCollection services = new ServiceCollection();
+        IDicomServerBuilder builder = Substitute.For<IDicomServerBuilder>();
+        builder.Services.Returns(services);
+
+        IConfiguration configuration = new ConfigurationBuilder().Build();
+
+        builder.AddSqlServer(configuration);
+
+        Assert.IsType<ActiveDirectoryAuthenticationProvider>(SqlAuthenticationProvider.GetProvider(SqlAuthenticationMethod.ActiveDirectoryManagedIdentity));
+        Assert.IsType<ActiveDirectoryAuthenticationProvider>(SqlAuthenticationProvider.GetProvider(SqlAuthenticationMethod.ActiveDirectoryMSI));
+    }
+
+    [Fact]
+    public void GivenFunctionsBuilder_WhenUsingDefaults_ThenUseBuiltInAuthenticationProvider()
+    {
+        IServiceCollection services = new ServiceCollection();
+        IDicomFunctionsBuilder builder = Substitute.For<IDicomFunctionsBuilder>();
+        builder.Services.Returns(services);
+
+        IConfiguration configuration = new ConfigurationBuilder()
+            .Build();
+
+        builder.AddSqlServer(configuration.Bind);
+
+        Assert.IsType<ActiveDirectoryAuthenticationProvider>(SqlAuthenticationProvider.GetProvider(SqlAuthenticationMethod.ActiveDirectoryManagedIdentity));
+        Assert.IsType<ActiveDirectoryAuthenticationProvider>(SqlAuthenticationProvider.GetProvider(SqlAuthenticationMethod.ActiveDirectoryMSI));
+    }
+
+    [Fact]
+    public void GivenWebServerBuilder_WhenConfiguringWorkloadIdenity_ThenIncludeAuthenticationProvider()
+    {
+        IServiceCollection services = new ServiceCollection();
+        IDicomServerBuilder builder = Substitute.For<IDicomServerBuilder>();
+        builder.Services.Returns(services);
+
+        IConfiguration configuration = new ConfigurationBuilder()
+            .AddInMemoryCollection(new KeyValuePair<string, string>[]
+            {
+                new KeyValuePair<string, string>($"{SqlServerDataStoreConfiguration.SectionName}:EnableWorkloadIdentity", "true"),
+            })
+            .Build();
+
+        builder.AddSqlServer(configuration);
+
+        Assert.IsNotType<ActiveDirectoryAuthenticationProvider>(SqlAuthenticationProvider.GetProvider(SqlAuthenticationMethod.ActiveDirectoryManagedIdentity));
+        Assert.IsNotType<ActiveDirectoryAuthenticationProvider>(SqlAuthenticationProvider.GetProvider(SqlAuthenticationMethod.ActiveDirectoryMSI));
+    }
+
+    [Fact]
+    public void GivenFunctionBuilder_WhenConfiguringWorkloadIdenity_ThenIncludeAuthenticationProvider()
+    {
+        IServiceCollection services = new ServiceCollection();
+        IDicomFunctionsBuilder builder = Substitute.For<IDicomFunctionsBuilder>();
+        builder.Services.Returns(services);
+
+        IConfiguration configuration = new ConfigurationBuilder()
+            .AddInMemoryCollection(new KeyValuePair<string, string>[]
+            {
+                new KeyValuePair<string, string>("EnableWorkloadIdentity", "true"),
+            })
+            .Build();
+
+        builder.AddSqlServer(configuration.Bind);
+
+        Assert.IsNotType<ActiveDirectoryAuthenticationProvider>(SqlAuthenticationProvider.GetProvider(SqlAuthenticationMethod.ActiveDirectoryManagedIdentity));
+        Assert.IsNotType<ActiveDirectoryAuthenticationProvider>(SqlAuthenticationProvider.GetProvider(SqlAuthenticationMethod.ActiveDirectoryMSI));
+    }
+
+    public void Dispose()
+    {
+        ActiveDirectoryAuthenticationProvider provider = new();
+        SqlAuthenticationProvider.SetProvider(SqlAuthenticationMethod.ActiveDirectoryManagedIdentity, provider);
+        SqlAuthenticationProvider.SetProvider(SqlAuthenticationMethod.ActiveDirectoryMSI, provider);
+    }
+}

src/Microsoft.Health.Dicom.SqlServer/Registration/DicomSqlServerRegistrationExtensions.cs

@@ -40,28 +40,18 @@ public static IDicomServerBuilder AddSqlServer(
         IConfiguration configurationRoot,
         Action<SqlServerDataStoreConfiguration> configureAction = null)
     {
-        IServiceCollection services = EnsureArg.IsNotNull(dicomServerBuilder, nameof(dicomServerBuilder)).Services;
-
-        // Add core SQL services
-        services
-            .AddSqlServerConnection(
-                config =>
-                {
-                    configurationRoot?.GetSection(SqlServerDataStoreConfiguration.SectionName).Bind(config);
-                    configureAction?.Invoke(config);
-                })
+        EnsureArg.IsNotNull(dicomServerBuilder, nameof(dicomServerBuilder));
+        EnsureArg.IsNotNull(configurationRoot, nameof(configurationRoot));
+
+        dicomServerBuilder.Services
+            .AddCommonSqlServices(sqlOptions =>
+            {
+                configurationRoot.GetSection(SqlServerDataStoreConfiguration.SectionName).Bind(sqlOptions);
+                configureAction?.Invoke(sqlOptions);
+            })
             .AddSqlServerManagement<SchemaVersion>()
             .AddSqlServerApi()
-            .AddBackgroundSqlSchemaVersionResolver();
-
-        // Add SQL-specific implementations
-        services
-            .AddSqlChangeFeedStores()
-            .AddSqlExtendedQueryTagStores()
-            .AddSqlExtendedQueryTagErrorStores()
-            .AddSqlIndexDataStores()
-            .AddSqlInstanceStores()
-            .AddSqlPartitionStores()
+            .AddBackgroundSqlSchemaVersionResolver()
             .AddSqlQueryStores()
             .AddSqlWorkitemStores();
 
@@ -75,14 +65,25 @@ public static IDicomFunctionsBuilder AddSqlServer(
         EnsureArg.IsNotNull(dicomFunctionsBuilder, nameof(dicomFunctionsBuilder));
         EnsureArg.IsNotNull(configureAction, nameof(configureAction));
 
-        IServiceCollection services = dicomFunctionsBuilder.Services;
+        dicomFunctionsBuilder.Services
+            .AddCommonSqlServices(configureAction)
+            .AddForegroundSqlSchemaVersionResolver();
 
+        return dicomFunctionsBuilder;
+    }
+
+    private static IServiceCollection AddCommonSqlServices(this IServiceCollection services, Action<SqlServerDataStoreConfiguration> configureAction)
+    {
         // Add core SQL services
-        services
-            .AddSqlServerConnection(configureAction)
-            .AddForegroundSqlSchemaVersionResolver();
+        services.AddSqlServerConnection(configureAction);
+
+        // Optionally enable workload identity
+        DicomSqlServerOptions options = new();
+        configureAction(options);
+        if (options.EnableWorkloadIdentity)
+            services.EnableWorkloadManagedIdentity();
 
-        // Add SQL-specific implementations
+        // Add SQL-specific data store implementations
         services
             .AddSqlExtendedQueryTagStores()
             .AddSqlExtendedQueryTagErrorStores()
@@ -91,7 +92,7 @@ public static IDicomFunctionsBuilder AddSqlServer(
             .AddSqlPartitionStores()
             .AddSqlChangeFeedStores();
 
-        return dicomFunctionsBuilder;
+        return services;
     }
 
     private static IServiceCollection AddBackgroundSqlSchemaVersionResolver(this IServiceCollection services)
@@ -203,4 +204,9 @@ private static IServiceCollection AddSqlWorkitemStores(this IServiceCollection s
 
         return services;
     }
+
+    private sealed class DicomSqlServerOptions : SqlServerDataStoreConfiguration
+    {
+        public bool EnableWorkloadIdentity { get; set; }
+    }
 }