Add interface for cosmosdb MI

src/Microsoft.Health.Fhir.Azure/ExportDestinationClient/AzureAccessTokenProvider.cs

@@ -14,7 +14,7 @@
 
 namespace Microsoft.Health.Fhir.Azure.ExportDestinationClient
 {
-    public class AzureAccessTokenProvider : IAccessTokenProvider
+    public class AzureAccessTokenProvider : IAccessTokenProvider, ICosmosDBAccessTokenProvider
     {
         private readonly DefaultAzureCredential _azureServiceTokenProvider;
         private readonly ILogger<AzureAccessTokenProvider> _logger;

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

@@ -126,5 +126,16 @@ public static IFhirServerBuilder AddAzureIntegrationDataStoreClient(this IFhirSe
 
             return fhirServerBuilder;
         }
+
+        public static IFhirServerBuilder AddCosmosdbMI(this IFhirServerBuilder fhirServerBuilder)
+        {
+            EnsureArg.IsNotNull(fhirServerBuilder, nameof(fhirServerBuilder));
+
+            fhirServerBuilder.Services.Add<AzureAccessTokenProvider>()
+                .Transient()
+                .AsService<ICosmosDBAccessTokenProvider>();
+
+            return fhirServerBuilder;
+        }
     }
 }

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

@@ -0,0 +1,23 @@
+// -------------------------------------------------------------------------------------------------
+// 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.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Health.Fhir.Core.Features.Operations
+{
+    public interface ICosmosDBAccessTokenProvider
+    {
+        /// <summary>
+        /// Gets the access token for the resource.
+        /// </summary>
+        /// <param name="resourceUri">Uri pointing to the resource for which we need an access token.</param>
+        /// <param name="cancellationToken">Cancellation token.</param>
+        /// <returns>Access token.</returns>
+        /// <exception cref="AccessTokenProviderException">Thrown when unable to get access token.</exception>
+        Task<string> GetAccessTokenForResourceAsync(Uri resourceUri, CancellationToken cancellationToken);
+    }
+}

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

@@ -32,7 +32,7 @@ internal class CosmosDbCollectionPhysicalPartitionInfo : IRequireInitializationO
         private readonly CosmosDataStoreConfiguration _dataStoreConfiguration;
         private readonly CosmosCollectionConfiguration _collectionConfiguration;
         private readonly IHttpClientFactory _httpClientFactory;
-        private readonly IAccessTokenProvider _accessTokenProvider;
+        private readonly ICosmosDBAccessTokenProvider _accessTokenProvider;
         private readonly ILogger<CosmosDbCollectionPhysicalPartitionInfo> _logger;
         private readonly CancellationTokenSource _backgroundLoopCancellationTokenSource = new();
         private Task _backgroundLoopTask;
@@ -41,7 +41,7 @@ public CosmosDbCollectionPhysicalPartitionInfo(
             CosmosDataStoreConfiguration dataStoreConfiguration,
             IOptionsMonitor<CosmosCollectionConfiguration> collectionConfiguration,
             IHttpClientFactory httpClientFactory,
-            IAccessTokenProvider accessTokenProvider,
+            ICosmosDBAccessTokenProvider accessTokenProvider,
             ILogger<CosmosDbCollectionPhysicalPartitionInfo> logger)
         {
             EnsureArg.IsNotNull(dataStoreConfiguration, nameof(dataStoreConfiguration));

src/Microsoft.Health.Fhir.Shared.Web/Startup.cs

@@ -121,6 +121,7 @@ private void AddDataStore(IServiceCollection services, IFhirServerBuilder fhirSe
             if (runtimeConfiguration is AzureApiForFhirRuntimeConfiguration)
             {
                 fhirServerBuilder.AddCosmosDb();
+                fhirServerBuilder.AddCosmosdbMI();
             }
             else if (runtimeConfiguration is AzureHealthDataServicesRuntimeConfiguration)
             {