Add interface for cosmosdb MI

src/Microsoft.Health.Fhir.Azure/FhirServerBuilderAzureRegistrationExtensions.cs

@@ -51,10 +51,6 @@ public static IFhirServerBuilder AddAzureExportClientInitializer(this IFhirServe
                 fhirServerBuilder.Services.Add<AzureAccessTokenClientInitializer>()
                     .Transient()
                     .AsService<IExportClientInitializer<BlobServiceClient>>();
-
-                fhirServerBuilder.Services.Add<AzureAccessTokenProvider>()
-                    .Transient()
-                    .AsService<IAccessTokenProvider>();
             }
             else
             {

src/Microsoft.Health.Fhir.Azure/ExportDestinationClient/AzureAccessTokenProvider.cs → src/Microsoft.Health.Fhir.Core/Features/Operations/AzureAccessTokenProvider.cs

@@ -4,15 +4,15 @@
 // -------------------------------------------------------------------------------------------------
 
 using System;
+using System.Globalization;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
 using Azure.Identity;
 using EnsureThat;
 using Microsoft.Extensions.Logging;
-using Microsoft.Health.Fhir.Core.Features.Operations;
 
-namespace Microsoft.Health.Fhir.Azure.ExportDestinationClient
+namespace Microsoft.Health.Fhir.Core.Features.Operations
 {
     public class AzureAccessTokenProvider : IAccessTokenProvider
     {
@@ -27,12 +27,14 @@ public AzureAccessTokenProvider(ILogger<AzureAccessTokenProvider> logger)
             _logger = logger;
         }
 
+        public TokenCredential TokenCredential => _azureServiceTokenProvider;
+
         public async Task<string> GetAccessTokenForResourceAsync(Uri resourceUri, CancellationToken cancellationToken)
         {
             EnsureArg.IsNotNull(resourceUri, nameof(resourceUri));
 
             // https://learn.microsoft.com/en-us/dotnet/api/overview/azure/app-auth-migration?view=azure-dotnet
-            var accessTokenContext = new TokenRequestContext(scopes: new[] { resourceUri + "/.default" });
+            var accessTokenContext = new TokenRequestContext(scopes: [resourceUri + "/.default"]);
 
             AccessToken accessToken;
             try
@@ -42,14 +44,14 @@ public async Task<string> GetAccessTokenForResourceAsync(Uri resourceUri, Cancel
             catch (CredentialUnavailableException ex)
             {
                 _logger.LogWarning(ex, "Failed to retrieve access token");
-                throw new AccessTokenProviderException(Resources.CannotGetAccessToken);
+                throw new AccessTokenProviderException(string.Format(CultureInfo.InvariantCulture, Core.Resources.CannotGetAccessToken, resourceUri));
             }
 
             if (string.IsNullOrEmpty(accessToken.Token))
             {
                 _logger.LogWarning("Failed to retrieve access token");
 
-                throw new AccessTokenProviderException(Resources.CannotGetAccessToken);
+                throw new AccessTokenProviderException(string.Format(CultureInfo.InvariantCulture, Core.Resources.CannotGetAccessToken, resourceUri));
             }
             else
             {

src/Microsoft.Health.Fhir.Core/Features/Operations/IAccessTokenProvider.cs

@@ -6,11 +6,14 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Azure.Core;
 
 namespace Microsoft.Health.Fhir.Core.Features.Operations
 {
     public interface IAccessTokenProvider
     {
+        TokenCredential TokenCredential { get; }
+
         /// <summary>
         /// Gets the access token for the resource.
         /// </summary>

src/Microsoft.Health.Fhir.Core/Resources.resx

@@ -744,4 +744,7 @@
     <value>The FHIR Server ran out of memory while processing an export job. Please use the _maxCount parameter when requesting an export job to reduce the number of resources exported at one time. The count used in this job was {0}</value>
     <comment>'{0}' the value of _maxCount used for this export job.</comment>
   </data>
+  <data name="CannotGetAccessToken" xml:space="preserve">
+    <value>Unable to get an access token for '{0}'.</value>
+  </data>
 </root>

src/Microsoft.Health.Fhir.CosmosDb.UnitTests/Features/Storage/FhirCosmosClientInitializerTests.cs

@@ -9,6 +9,7 @@
 using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Health.Core.Features.Context;
 using Microsoft.Health.Fhir.Core.Features.Context;
+using Microsoft.Health.Fhir.Core.Features.Operations;
 using Microsoft.Health.Fhir.CosmosDb.Core.Configs;
 using Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage;
 using Microsoft.Health.Fhir.CosmosDb.Features.Storage;
@@ -40,7 +41,7 @@ public FhirCosmosClientInitializerTests()
                 clientTestProvider,
                 () => new[] { new TestRequestHandler() },
                 new RetryExceptionPolicyFactory(_cosmosDataStoreConfiguration, Substitute.For<RequestContextAccessor<IFhirRequestContext>>()),
-                new Lazy<TokenCredential>(() => Substitute.For<TokenCredential>()),
+                Substitute.For<CosmosAccessTokenProviderFactory>(),
                 NullLogger<FhirCosmosClientInitializer>.Instance);
 
             _collectionInitializers = new List<ICollectionInitializer> { _collectionInitializer1, _collectionInitializer2 };

src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/CosmosAccessTokenProviderFactory.cs

@@ -0,0 +1,13 @@
+// -------------------------------------------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+// -------------------------------------------------------------------------------------------------
+
+using Microsoft.Health.Fhir.Core.Features.Operations;
+
+namespace Microsoft.Health.Fhir.CosmosDb.Features.Storage;
+
+/// <summary>
+/// Factory for creating <see cref="IAccessTokenProvider"/> instances for access to Cosmos.
+/// </summary>
+public delegate IAccessTokenProvider CosmosAccessTokenProviderFactory();

src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/CosmosDbCollectionPhysicalPartitionInfo.cs

@@ -41,7 +41,7 @@ public CosmosDbCollectionPhysicalPartitionInfo(
             CosmosDataStoreConfiguration dataStoreConfiguration,
             IOptionsMonitor<CosmosCollectionConfiguration> collectionConfiguration,
             IHttpClientFactory httpClientFactory,
-            IAccessTokenProvider accessTokenProvider,
+            CosmosAccessTokenProviderFactory accessTokenProviderFactory,
             ILogger<CosmosDbCollectionPhysicalPartitionInfo> logger)
         {
             EnsureArg.IsNotNull(dataStoreConfiguration, nameof(dataStoreConfiguration));
@@ -52,7 +52,7 @@ public CosmosDbCollectionPhysicalPartitionInfo(
             _dataStoreConfiguration = dataStoreConfiguration;
             _collectionConfiguration = collectionConfiguration.Get(Constants.CollectionConfigurationName);
             _httpClientFactory = httpClientFactory;
-            _accessTokenProvider = accessTokenProvider;
+            _accessTokenProvider = accessTokenProviderFactory.Invoke();
             _logger = logger;
         }
 

src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/FhirCosmosClientInitializer.cs

@@ -14,6 +14,7 @@
 using Microsoft.Azure.Cosmos.Fluent;
 using Microsoft.Extensions.Logging;
 using Microsoft.Health.Abstractions.Exceptions;
+using Microsoft.Health.Fhir.Core.Features.Operations;
 using Microsoft.Health.Fhir.CosmosDb.Core.Configs;
 using Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage;
 using Microsoft.IO;
@@ -27,11 +28,11 @@ public class FhirCosmosClientInitializer : ICosmosClientInitializer
         private const StringComparison _hashCodeStringComparison = StringComparison.Ordinal;
 
         private readonly ICosmosClientTestProvider _testProvider;
-        private readonly ILogger<FhirCosmosClientInitializer> _logger;
         private readonly Func<IEnumerable<RequestHandler>> _requestHandlerFactory;
         private readonly RetryExceptionPolicyFactory _retryExceptionPolicyFactory;
-        private readonly Lazy<TokenCredential> _tokenCredential;
+        private readonly CosmosAccessTokenProviderFactory _cosmosAccessTokenProviderFactory;
         private readonly object _lockObject;
+        private readonly ILogger<FhirCosmosClientInitializer> _logger;
 
         private CosmosClient _cosmosClient;
         private int _cosmosKeyHashCode;
@@ -40,19 +41,19 @@ public FhirCosmosClientInitializer(
             ICosmosClientTestProvider testProvider,
             Func<IEnumerable<RequestHandler>> requestHandlerFactory,
             RetryExceptionPolicyFactory retryExceptionPolicyFactory,
-            Lazy<TokenCredential> tokenCredential,
+            CosmosAccessTokenProviderFactory cosmosAccessTokenProviderFactory,
             ILogger<FhirCosmosClientInitializer> logger)
         {
             EnsureArg.IsNotNull(testProvider, nameof(testProvider));
             EnsureArg.IsNotNull(requestHandlerFactory, nameof(requestHandlerFactory));
             EnsureArg.IsNotNull(retryExceptionPolicyFactory, nameof(retryExceptionPolicyFactory));
-            EnsureArg.IsNotNull(tokenCredential, nameof(tokenCredential));
+            EnsureArg.IsNotNull(cosmosAccessTokenProviderFactory, nameof(cosmosAccessTokenProviderFactory));
             EnsureArg.IsNotNull(logger, nameof(logger));
 
             _testProvider = testProvider;
             _requestHandlerFactory = requestHandlerFactory;
             _retryExceptionPolicyFactory = retryExceptionPolicyFactory;
-            _tokenCredential = tokenCredential;
+            _cosmosAccessTokenProviderFactory = cosmosAccessTokenProviderFactory;
             _logger = logger;
             _lockObject = new object();
 
@@ -127,7 +128,7 @@ private CosmosClient CreateCosmosClientInternal(CosmosDataStoreConfiguration con
             CosmosClientBuilder builder;
 
             builder = configuration.UseManagedIdentity ?
-                new CosmosClientBuilder(host, _tokenCredential.Value) :
+                new CosmosClientBuilder(host, _cosmosAccessTokenProviderFactory.Invoke().TokenCredential) :
                 new CosmosClientBuilder(host, key);
 
             builder

src/Microsoft.Health.Fhir.CosmosDb/Registration/FhirServerBuilderCosmosDbRegistrationExtensions.cs

@@ -15,6 +15,7 @@
 using Microsoft.Extensions.Options;
 using Microsoft.Health.Extensions.DependencyInjection;
 using Microsoft.Health.Fhir.Core.Extensions;
+using Microsoft.Health.Fhir.Core.Features.Operations;
 using Microsoft.Health.Fhir.Core.Features.Operations.Export;
 using Microsoft.Health.Fhir.Core.Features.Parameters;
 using Microsoft.Health.Fhir.Core.Features.Search.Expressions;
@@ -271,6 +272,10 @@ private static IFhirServerBuilder AddCosmosDbPersistence(this IFhirServerBuilder
                 .AsImplementedInterfaces()
                 .AsFactory<IScoped<IQueueClient>>();
 
+            services.Add<CosmosAccessTokenProviderFactory>(c => c.GetRequiredService<IAccessTokenProvider>)
+               .Transient()
+               .AsSelf();
+
             IEnumerable<TypeRegistrationBuilder> jobs = services.TypesInSameAssemblyAs<CosmosExportOrchestratorJob>()
                 .AssignableTo<IJob>()
                 .Transient()

src/Microsoft.Health.Fhir.Shared.Api/Modules/OperationsModule.cs

@@ -10,6 +10,7 @@
 using Microsoft.Health.Fhir.Api.Features.Operations;
 using Microsoft.Health.Fhir.Core.Extensions;
 using Microsoft.Health.Fhir.Core.Features.Conformance;
+using Microsoft.Health.Fhir.Core.Features.Operations;
 using Microsoft.Health.Fhir.Core.Features.Operations.Everything;
 using Microsoft.Health.Fhir.Core.Features.Operations.Export;
 using Microsoft.Health.Fhir.Core.Features.Operations.Import;
@@ -88,6 +89,10 @@ public void Load(IServiceCollection services)
                 .Transient()
                 .AsSelf()
                 .AsImplementedInterfaces();
+
+            services.Add<AzureAccessTokenProvider>()
+                .Transient()
+                .AsService<IAccessTokenProvider>();
         }
     }
 }

test/Microsoft.Health.Fhir.Shared.Tests.Integration/Persistence/CosmosDbFhirStorageTestsFixture.cs

@@ -142,14 +142,17 @@ public virtual async Task InitializeAsync()
 
             var cosmosResponseProcessor = Substitute.For<ICosmosResponseProcessor>();
 
+            var cosmosAccessor = Substitute.For<IAccessTokenProvider>();
+            cosmosAccessor.TokenCredential.Returns(GetTokenCredential());
+
             var responseProcessor = new CosmosResponseProcessor(_fhirRequestContextAccessor, mediator, Substitute.For<ICosmosQueryLogger>(), NullLogger<CosmosResponseProcessor>.Instance);
             var handler = new FhirCosmosResponseHandler(() => new NonDisposingScope(_container), _cosmosDataStoreConfiguration, _fhirRequestContextAccessor, responseProcessor);
             var retryExceptionPolicyFactory = new RetryExceptionPolicyFactory(_cosmosDataStoreConfiguration, _fhirRequestContextAccessor);
             var documentClientInitializer = new FhirCosmosClientInitializer(
                 testProvider,
                 () => new[] { handler },
                 retryExceptionPolicyFactory,
-                new Lazy<TokenCredential>(GetTokenCredential),
+                () => cosmosAccessor,
                 NullLogger<FhirCosmosClientInitializer>.Instance);
             _cosmosClient = documentClientInitializer.CreateCosmosClient(_cosmosDataStoreConfiguration);