Integration tests use DevOps MI for Cosmos

build/jobs/provision-deploy.yml

@@ -128,6 +128,28 @@ jobs:
             }
           }
         }
+
+        if("${{ parameters.sql }}" -eq "false"){
+            Write-Host "Add DevOps MI permission to Cosmos databases"
+
+            $principalId = (Get-AzContext).Account.Id
+
+            New-AzCosmosDBSqlRoleAssignment `
+              -AccountName $webAppName `
+              -ResourceGroupName $resourceGroupName `
+              -Scope "/" `
+              -PrincipalId $principalId `
+              -RoleDefinitionId "00000000-0000-0000-0000-000000000002"
+
+            $roleDefinitionId = "230815da-be43-4aae-9cb4-875f7bd000aa"
+            $scope = "/subscriptions/$((Get-AzContext).Subscription.Id)/resourceGroups/$($resourceGroupName)/providers/Microsoft.DocumentDB/databaseAccounts/$($webAppName)"
+
+            New-AzRoleAssignment `
+                -ObjectId $principalId `
+                -RoleDefinitionId $roleDefinitionId `
+                -Scope $scope
+        }
+
   - template: ./provision-healthcheck.yml
     parameters: 
       webAppName: ${{ parameters.webAppName }}

build/jobs/run-tests.yml

@@ -36,6 +36,24 @@ jobs:
     inputs:
       azureSubscription: $(ConnectedServiceName)
       KeyVaultName: '${{ parameters.keyVaultName }}'
+
+  - task: AzurePowerShell@5
+    displayName: 'Set Workload Identity Variables'
+    inputs:
+      azureSubscription: $(ConnectedServiceName)
+      azurePowerShellVersion: latestVersion
+      ScriptType: inlineScript
+      Inline: |
+        Write-Host "##vso[task.setvariable variable=AZURESUBSCRIPTION_CLIENT_ID]$env:AZURESUBSCRIPTION_CLIENT_ID"
+        Write-Host "##vso[task.setvariable variable=AZURESUBSCRIPTION_TENANT_ID]$env:AZURESUBSCRIPTION_TENANT_ID"
+        Write-Host "##vso[task.setvariable variable=AZURESUBSCRIPTION_SERVICE_CONNECTION_ID]$env:AZURESUBSCRIPTION_SERVICE_CONNECTION_ID"
+
+        $appSettings = (Get-AzWebApp -ResourceGroupName $(ResourceGroupName) -Name $appServiceName).SiteConfig.AppSettings
+        $dataStoreResourceId = $appSettings | where {$_.Name -eq "FhirServer__ResourceManager__DataStoreResourceId"}
+        $dataStoreResourceId = $dataStoreResourceId[0].Value
+        Write-Host "$dataStoreResourceId"
+        Write-Host "##vso[task.setvariable variable=DataStoreResourceId]$($dataStoreResourceId)"
+
 
   - task: DotNetCoreCLI@2
     displayName: 'Run Cosmos Integration Tests'
@@ -46,7 +64,11 @@ jobs:
       testRunTitle: '${{ parameters.version }} Integration'
     env:
       'CosmosDb:Host': $(CosmosDb--Host)
-      'CosmosDb:Key': $(CosmosDb--Key)
+      'AZURESUBSCRIPTION_CLIENT_ID': '$(AZURESUBSCRIPTION_CLIENT_ID)'
+      'AZURESUBSCRIPTION_TENANT_ID': '$(AZURESUBSCRIPTION_TENANT_ID)'
+      'AZURESUBSCRIPTION_SERVICE_CONNECTION_ID': '$(AZURESUBSCRIPTION_SERVICE_CONNECTION_ID)'
+      'SYSTEM_ACCESSTOKEN': $(System.AccessToken)
+      'FhirServer:ResourceManager:DataStoreResourceId': $(DataStoreResourceId)
 
 - job: "SqlIntegrationTests"
   pool:

src/Microsoft.Health.Fhir.CosmosDb.Initialization/Features/Storage/ResourceManagerCollectionSetup.cs

@@ -46,6 +46,7 @@ public ResourceManagerCollectionSetup(
         CosmosDataStoreConfiguration cosmosDataStoreConfiguration,
         IConfiguration genericConfiguration,
         IEnumerable<IStoredProcedureMetadata> storedProcedures,
+        TokenCredential tokenCredential,
         ILogger<ResourceManagerCollectionSetup> logger)
     {
         EnsureArg.IsNotNull(storedProcedures, nameof(storedProcedures));
@@ -54,8 +55,9 @@ public ResourceManagerCollectionSetup(
         _cosmosDataStoreConfiguration = EnsureArg.IsNotNull(cosmosDataStoreConfiguration, nameof(cosmosDataStoreConfiguration));
         _storeProceduresMetadata = EnsureArg.IsNotNull(storedProcedures, nameof(storedProcedures));
         _logger = EnsureArg.IsNotNull(logger, nameof(logger));
+        EnsureArg.IsNotNull(tokenCredential, nameof(tokenCredential));
 
-        _armClient = new ArmClient(new DefaultAzureCredential());
+        _armClient = new ArmClient(tokenCredential);
         var dataStoreResourceId = genericConfiguration.GetValue("FhirServer:ResourceManager:DataStoreResourceId", string.Empty);
         _resourceIdentifier = ResourceIdentifier.Parse(dataStoreResourceId);
         _account = _armClient.GetCosmosDBAccountResource(_resourceIdentifier);

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

@@ -5,6 +5,7 @@
 
 using System;
 using System.Collections.Generic;
+using Azure.Core;
 using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Health.Core.Features.Context;
 using Microsoft.Health.Fhir.Core.Features.Context;
@@ -39,6 +40,7 @@ public FhirCosmosClientInitializerTests()
                 clientTestProvider,
                 () => new[] { new TestRequestHandler() },
                 new RetryExceptionPolicyFactory(_cosmosDataStoreConfiguration, Substitute.For<RequestContextAccessor<IFhirRequestContext>>()),
+                new Lazy<TokenCredential>(() => Substitute.For<TokenCredential>()),
                 NullLogger<FhirCosmosClientInitializer>.Instance);
 
             _collectionInitializers = new List<ICollectionInitializer> { _collectionInitializer1, _collectionInitializer2 };

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

@@ -8,7 +8,7 @@
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
-using Azure.Identity;
+using Azure.Core;
 using EnsureThat;
 using Microsoft.Azure.Cosmos;
 using Microsoft.Azure.Cosmos.Fluent;
@@ -30,6 +30,7 @@ public class FhirCosmosClientInitializer : ICosmosClientInitializer
         private readonly ILogger<FhirCosmosClientInitializer> _logger;
         private readonly Func<IEnumerable<RequestHandler>> _requestHandlerFactory;
         private readonly RetryExceptionPolicyFactory _retryExceptionPolicyFactory;
+        private readonly Lazy<TokenCredential> _tokenCredential;
         private readonly object _lockObject;
 
         private CosmosClient _cosmosClient;
@@ -39,16 +40,19 @@ public FhirCosmosClientInitializer(
             ICosmosClientTestProvider testProvider,
             Func<IEnumerable<RequestHandler>> requestHandlerFactory,
             RetryExceptionPolicyFactory retryExceptionPolicyFactory,
+            Lazy<TokenCredential> tokenCredential,
             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(logger, nameof(logger));
 
             _testProvider = testProvider;
             _requestHandlerFactory = requestHandlerFactory;
             _retryExceptionPolicyFactory = retryExceptionPolicyFactory;
+            _tokenCredential = tokenCredential;
             _logger = logger;
             _lockObject = new object();
 
@@ -123,7 +127,7 @@ private CosmosClient CreateCosmosClientInternal(CosmosDataStoreConfiguration con
             CosmosClientBuilder builder;
 
             builder = configuration.UseManagedIdentity ?
-                new CosmosClientBuilder(host, new DefaultAzureCredential()) :
+                new CosmosClientBuilder(host, _tokenCredential.Value) :
                 new CosmosClientBuilder(host, key);
 
             builder

src/Microsoft.Health.Fhir.Shared.Api/Microsoft.Health.Fhir.Shared.Api.projitems

@@ -53,6 +53,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)Features\Routing\QueryStringParser.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Modules\FeatureFlags\Validate\ValidateFeatureModule.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Modules\FeatureFlags\Validate\ValidatePostConfigureOptions.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)Modules\IdentityModule.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Modules\MediationModule.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Modules\MvcModule.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Features\Exceptions\BaseExceptionMiddleware.cs" />

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

@@ -0,0 +1,21 @@
+// -------------------------------------------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+// -------------------------------------------------------------------------------------------------
+
+using Azure.Core;
+using Azure.Identity;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection.Extensions;
+using Microsoft.Health.Extensions.DependencyInjection;
+
+namespace Microsoft.Health.Fhir.Api.Modules;
+
+public class IdentityModule : IStartupModule
+{
+    public void Load(IServiceCollection services)
+    {
+        // Provide support for Azure Managed Identity
+        services.TryAddSingleton<TokenCredential, DefaultAzureCredential>();
+    }
+}

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

@@ -12,7 +12,6 @@
 using Microsoft.Health.Extensions.DependencyInjection;
 using Microsoft.Health.Fhir.Api.Features.Binders;
 using Microsoft.Health.Fhir.Api.Features.Filters;
-using Microsoft.Health.Fhir.Api.Features.Resources.Bundle;
 using Microsoft.Health.Fhir.Api.Features.Routing;
 using Microsoft.Health.Fhir.Core.Features.Routing;
 

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

@@ -11,7 +11,6 @@
 using Microsoft.Health.Extensions.DependencyInjection;
 using Microsoft.Health.Fhir.Api.Configs;
 using Microsoft.Health.Fhir.Api.Features.Filters;
-using Microsoft.Health.Fhir.Api.Features.Resources.Bundle;
 using Microsoft.Health.Fhir.Api.Features.Routing;
 using Microsoft.Health.Fhir.Core.Extensions;
 using Microsoft.Health.Fhir.Core.Features.Compartment;

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

@@ -8,6 +8,8 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Azure.Core;
+using Azure.Identity;
 using MediatR;
 using Microsoft.Azure.Cosmos;
 using Microsoft.Extensions.Logging.Abstractions;
@@ -82,7 +84,9 @@ public CosmosDbFhirStorageTestsFixture()
                 Host = Environment.GetEnvironmentVariable("CosmosDb:Host") ?? CosmosDbLocalEmulator.Host,
                 Key = Environment.GetEnvironmentVariable("CosmosDb:Key") ?? CosmosDbLocalEmulator.Key,
                 DatabaseId = Environment.GetEnvironmentVariable("CosmosDb:DatabaseId") ?? "FhirTests",
+                UseManagedIdentity = bool.TryParse(Environment.GetEnvironmentVariable("CosmosDb:UseManagedIdentity"), out bool useManagedIdentity) && useManagedIdentity,
                 AllowDatabaseCreation = true,
+                AllowCollectionSetup = true,
                 PreferredLocations = Environment.GetEnvironmentVariable("CosmosDb:PreferredLocations")?.Split(';', StringSplitOptions.RemoveEmptyEntries),
                 UseQueueClientJobs = true,
             };
@@ -139,22 +143,39 @@ public virtual async Task InitializeAsync()
             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, NullLogger<FhirCosmosClientInitializer>.Instance);
+            var documentClientInitializer = new FhirCosmosClientInitializer(
+                testProvider,
+                () => new[] { handler },
+                retryExceptionPolicyFactory,
+                new Lazy<TokenCredential>(GetTokenCredential),
+                NullLogger<FhirCosmosClientInitializer>.Instance);
             _cosmosClient = documentClientInitializer.CreateCosmosClient(_cosmosDataStoreConfiguration);
 
             var fhirCollectionInitializer = new CollectionInitializer(_cosmosCollectionConfiguration, _cosmosDataStoreConfiguration, upgradeManager, testProvider, NullLogger<CollectionInitializer>.Instance);
 
             // Cosmos DB emulators throws errors when multiple collections are initialized concurrently.
             // Use the semaphore to only allow one initialization at a time.
             await CollectionInitializationSemaphore.WaitAsync();
-            var dataCollectionSetup = new DataPlaneCollectionSetup(
-                _cosmosDataStoreConfiguration,
-                optionsMonitor,
-                documentClientInitializer,
-                storedProcedureInstaller,
-                NullLogger<DataPlaneCollectionSetup>.Instance);
             try
             {
+                ICollectionSetup dataCollectionSetup;
+
+                if (_cosmosDataStoreConfiguration.UseManagedIdentity)
+                {
+                    dataCollectionSetup = new ResourceManagerCollectionSetup(
+                        new OptionsMonitor<CosmosCollectionConfiguration>(_cosmosCollectionConfiguration)
+                        )
+                }
+                else
+                {
+                    dataCollectionSetup = new DataPlaneCollectionSetup(
+                        _cosmosDataStoreConfiguration,
+                        optionsMonitor,
+                        documentClientInitializer,
+                        storedProcedureInstaller,
+                        NullLogger<DataPlaneCollectionSetup>.Instance);
+                }
+
                 await dataCollectionSetup.CreateDatabaseAsync(retryExceptionPolicyFactory.RetryPolicy, CancellationToken.None);
                 await dataCollectionSetup.CreateCollectionAsync(new List<ICollectionInitializer> { fhirCollectionInitializer }, retryExceptionPolicyFactory.RetryPolicy, CancellationToken.None);
                 await dataCollectionSetup.InstallStoredProcs(CancellationToken.None);
@@ -271,6 +292,23 @@ public virtual async Task DisposeAsync()
             _cosmosClient.Dispose();
         }
 
+        private TokenCredential GetTokenCredential()
+        {
+            // Add custom logic to set up the AzurePipelinesCredential if we are running in Azure Pipelines
+            string federatedClientId = Environment.GetEnvironmentVariable("AZURESUBSCRIPTION_CLIENT_ID");
+            string federatedTenantId = Environment.GetEnvironmentVariable("AZURESUBSCRIPTION_TENANT_ID");
+            string serviceConnectionId = Environment.GetEnvironmentVariable("AZURESUBSCRIPTION_SERVICE_CONNECTION_ID");
+            string systemAccessToken = Environment.GetEnvironmentVariable("SYSTEM_ACCESSTOKEN");
+
+            if (!string.IsNullOrEmpty(federatedClientId) && !string.IsNullOrEmpty(federatedTenantId) && !string.IsNullOrEmpty(serviceConnectionId) && !string.IsNullOrEmpty(systemAccessToken))
+            {
+                AzurePipelinesCredential azurePipelinesCredential = new(federatedTenantId, federatedClientId, serviceConnectionId, systemAccessToken);
+                return azurePipelinesCredential;
+            }
+
+            return new DefaultAzureCredential();
+        }
+
         object IServiceProvider.GetService(Type serviceType)
         {
             if (serviceType == typeof(IFhirDataStore))