[Feature] Changes Cosmos to use MI connection (#4026)

* Cosmos OSS MI * Update to install procs with ARM * Refactored cosmos initialization
microsoft · Sep 6, 2024 · 66ba9a5 · 66ba9a5
1 parent 14494b0
commit 66ba9a5
Show file tree

Hide file tree

Showing 29 changed files with 592 additions and 138 deletions.
diff --git a/Directory.Packages.props b/Directory.Packages.props
@@ -32,6 +32,7 @@
   <ItemGroup>
     <PackageVersion Include="AngleSharp" Version="1.1.2" />
     <PackageVersion Include="Azure.Identity" Version="1.12.0" />
+    <PackageVersion Include="Azure.ResourceManager.CosmosDB" Version="1.3.2" />
     <PackageVersion Include="Azure.Extensions.AspNetCore.Configuration.Secrets" Version="1.3.1" />
     <PackageVersion Include="Azure.Monitor.OpenTelemetry.AspNetCore" Version="1.2.0" />
     <PackageVersion Include="Azure.Storage.Blobs" Version="12.21.2" />

diff --git a/build/jobs/add-aad-test-environment.yml b/build/jobs/add-aad-test-environment.yml
@@ -42,12 +42,12 @@ steps:
 
       $vaultName = "${environmentName}-ts"
       $vaultLocation = "westus"
-      $vaultResourceGroupName = $ResourceGroupName
+      $vaultResourceGroupName = "$(ResourceGroupName)"
       if (Get-AzKeyVault -VaultName $vaultName -Location $vaultLocation -InRemovedState)
       {
           Write-Host "Attempting to restore vault ${vaultName}"
 
-          Undo-AzKeyVaultRemoval -VaultName $vaultName -Location $vaultLocation -Confirm
+          Undo-AzKeyVaultRemoval -VaultName $vaultName -ResourceGroupName $vaultResourceGroupName -Location $vaultLocation
           Write-Host "KeyVault $vaultName is restored"
       }
       Write-Host "Restored keyvaults"

diff --git a/build/jobs/provision-deploy.yml b/build/jobs/provision-deploy.yml
@@ -92,19 +92,19 @@ jobs:
             $templateParameters["sqlSchemaAutomaticUpdatesEnabled"] = "${{parameters.schemaAutomaticUpdatesEnabled}}"
         }
 
+        $deploymentName = $webAppName
+        $resourceGroupName = "${{ parameters.resourceGroup }}"
+
         Write-Host "Check for keyvaults in removed state..."
         if (Get-AzKeyVault -VaultName $webAppName -Location $(ResourceGroupRegion) -InRemovedState)
         {
-            Undo-AzKeyVaultRemoval -VaultName $webAppName -ResourceGroupName $parameters.resourceGroup -Location $(ResourceGroupRegion) -Confirm
+            Undo-AzKeyVaultRemoval -VaultName $webAppName -ResourceGroupName $resourceGroupName -Location $(ResourceGroupRegion)
             Write-Host "KeyVault $webAppName is restored"
         }
 
         Write-Host "Provisioning Resource Group"
         Write-Host "ResourceGroupName: ${{ parameters.resourceGroup }}"
 
-        $deploymentName = $webAppName
-        $resourceGroupName = "${{ parameters.resourceGroup }}"
-
         # Check if a deployment with the specified name already exists
         $existingDeployment = Get-AzResourceGroupDeployment -Name $deploymentName -ResourceGroupName $resourceGroupName -ErrorAction SilentlyContinue
 

diff --git a/samples/templates/default-azuredeploy-docker.json b/samples/templates/default-azuredeploy-docker.json
@@ -250,6 +250,8 @@
         "appInsightsName": "[concat('AppInsights-', variables('serviceName'))]",
         "storageBlobDataContributerRoleId": "[concat('/subscriptions/', subscription().subscriptionId, '/providers/Microsoft.Authorization/roleDefinitions/', 'ba92f5b4-2d11-453d-a403-e96b0029c9fe')]",
         "acrPullRoleId": "[concat('/subscriptions/', subscription().subscriptionId, '/providers/Microsoft.Authorization/roleDefinitions/', '7f951dda-4ed3-4680-a7ca-43fe172d538d')]",
+        "cosmosDbDataContributorRoleId": "[resourceId('Microsoft.DocumentDB/databaseAccounts/sqlRoleDefinitions', variables('serviceName'), '00000000-0000-0000-0000-000000000002')]",
+        "cosmosDbControlPlaneContributorRoleId": "[concat('/subscriptions/', subscription().subscriptionId, '/providers/Microsoft.Authorization/roleDefinitions/', '230815da-be43-4aae-9cb4-875f7bd000aa')]",
         "blobStorageUri": "[if(variables('isMAG'), '.blob.core.usgovcloudapi.net', '.blob.core.windows.net')]",
         "enableIntegrationStore": "[or(parameters('enableExport'), parameters('enableImport'))]",
         "staticFhirServerConfigProperties": {
@@ -262,6 +264,7 @@
             "FhirServer__Security__EnableAadSmartOnFhirProxy": "[parameters('enableAadSmartOnFhirProxy')]",
             "FhirServer__Security__Authentication__Authority": "[parameters('securityAuthenticationAuthority')]",
             "FhirServer__Security__Authentication__Audience": "[parameters('securityAuthenticationAudience')]",
+            "CosmosDb__UseManagedIdentity": "true",
             "CosmosDb__ContinuationTokenSizeLimitInKb": "1",
             "CosmosDb__UseQueueClientJobs": "true",
             "SqlServer__Initialize": "[equals(parameters('solutionType'),'FhirServerSqlServer')]",
@@ -364,7 +367,8 @@
             },
             "dependsOn": [
                 "[resourceId('Microsoft.Web/serverfarms', variables('appServicePlanName'))]",
-                "[resourceId('Microsoft.ManagedIdentity/userAssignedIdentities', variables('managedIdentityName'))]"
+                "[resourceId('Microsoft.ManagedIdentity/userAssignedIdentities', variables('managedIdentityName'))]",
+                "[if(variables('deployAppInsights'),concat('Microsoft.Insights/components/', variables('appInsightsName')),resourceId('Microsoft.KeyVault/vaults', variables('serviceName')))]"
             ],
             "resources": [
                 {
@@ -376,7 +380,7 @@
                         "[if(variables('deployAppInsights'),concat('Microsoft.Insights/components/', variables('appInsightsName')),resourceId('Microsoft.KeyVault/vaults', variables('serviceName')))]",
                         "[if(equals(parameters('solutionType'),'FhirServerCosmosDB'), resourceId('Microsoft.KeyVault/vaults/secrets', variables('serviceName'), 'CosmosDb--Host'), resourceId('Microsoft.KeyVault/vaults/secrets', variables('serviceName'), 'SqlServer--ConnectionString'))]"
                     ],
-                    "properties": "[if(variables('deployAppInsights'), union(variables('combinedFhirServerConfigProperties'), json(concat('{', '\"Telemetry__Provider\": \"', parameters('telemetryProviderType'), '\",', '\"Telemetry__InstrumentationKey\": \"', reference(concat('Microsoft.Insights/components/', variables('appInsightsName'))).InstrumentationKey, '\",', '\"Telemetry__ConnectionString\": \"', reference(concat('Microsoft.Insights/components/', variables('appInsightsName'))).ConnectionString, '\"', '}'))), variables('combinedFhirServerConfigProperties'))]"
+                    "properties": "[union(variables('combinedFhirServerConfigProperties'), json(concat('{ \"FhirServer__ResourceManager__DataStoreResourceId\": \"', if(equals(parameters('solutionType'),'FhirServerCosmosDB'), resourceId('Microsoft.DocumentDb/databaseAccounts', variables('serviceName')), resourceId('Microsoft.Sql/servers/', variables('sqlServerDerivedName'))), '\", ', if(variables('deployAppInsights'), concat('\"Telemetry__Provider\": \"', parameters('telemetryProviderType'), '\",', '\"Telemetry__InstrumentationKey\": \"', reference(resourceId('Microsoft.Insights/components', variables('appInsightsName'))).InstrumentationKey, '\",', '\"Telemetry__ConnectionString\": \"', reference(resourceId('Microsoft.Insights/components', variables('appInsightsName'))).ConnectionString, '\"'), ''), '}')))]"
                 },
                 {
                     "apiVersion": "2018-11-01",
@@ -461,6 +465,50 @@
                 ]
             }
         },
+        {
+            "condition": "[equals(parameters('solutionType'),'FhirServerCosmosDB')]",
+            "type": "Microsoft.Authorization/roleAssignments",
+            "apiVersion": "2020-04-01-preview",
+            "name": "[guid(uniqueString(variables('cosmosDbControlPlaneContributorRoleId'), parameters('fhirVersion'), variables('serviceName')))]",
+            "dependsOn": [
+                "[resourceId('Microsoft.Web/sites/', variables('serviceName'))]",
+                "[resourceId('Microsoft.DocumentDB/databaseAccounts', variables('serviceName'))]"
+            ],
+            "properties": {
+                "roleDefinitionId": "[variables('cosmosDbControlPlaneContributorRoleId')]",
+                "principalId": "[reference(resourceId('Microsoft.Web/sites/', variables('serviceName')), '2020-06-01', 'Full').identity.principalId]",
+                "principalType": "ServicePrincipal"
+            }
+        },
+        {
+            "condition": "[equals(parameters('solutionType'),'FhirServerCosmosDB')]",
+            "type": "Microsoft.DocumentDB/databaseAccounts/sqlDatabases",
+            "apiVersion": "2022-11-15",
+            "name": "[concat(variables('serviceName'), '/health')]",
+            "dependsOn": [
+                "[resourceId('Microsoft.DocumentDB/databaseAccounts', variables('serviceName'))]"
+            ],
+            "properties": {
+                "resource": {
+                    "id": "health"
+                }
+            }
+        },
+        {
+            "condition": "[equals(parameters('solutionType'),'FhirServerCosmosDB')]",
+            "apiVersion": "2022-11-15",
+            "name": "[concat(variables('serviceName'), '/', guid(uniqueString('CosmosDB', parameters('fhirVersion'), variables('serviceName'))))]",
+            "type": "Microsoft.DocumentDB/databaseAccounts/sqlRoleAssignments",
+            "dependsOn": [
+                "[resourceId('Microsoft.DocumentDB/databaseAccounts', variables('serviceName'))]",
+                "[variables('appServiceResourceId')]"
+            ],
+            "properties": {
+                "roleDefinitionId": "[variables('cosmosDbDataContributorRoleId')]",
+                "principalId": "[reference(concat('Microsoft.Web/sites/', variables('serviceName')), '2018-11-01', 'full').identity.principalId]",
+                "scope": "[resourceId('Microsoft.DocumentDb/databaseAccounts', variables('serviceName'))]"
+            }
+        },
         {
             "condition": "[and(equals(parameters('solutionType'),'FhirServerSqlServer'),equals(parameters('sqlServerNewOrExisting'), 'new'))]",
             "name": "[variables('sqlServerDerivedName')]",
@@ -514,6 +562,7 @@
             "type": "Microsoft.Resources/deployments",
             "apiVersion": "2022-09-01",
             "name": "webAppIPsSQLFirewall",
+            "condition": "[equals(parameters('solutionType'),'FhirServerSqlServer')]",
             "properties": {
                 "expressionEvaluationOptions": {
                     "scope": "inner"

diff --git a/src/Microsoft.Health.Fhir.Api/Configs/FhirServerConfiguration.cs b/src/Microsoft.Health.Fhir.Api/Configs/FhirServerConfiguration.cs
@@ -36,5 +36,7 @@ public class FhirServerConfiguration : IApiConfiguration
         public ImplementationGuidesConfiguration ImplementationGuides { get; } = new ImplementationGuidesConfiguration();
 
         public EncryptionConfiguration Encryption { get; } = new EncryptionConfiguration();
+
+        public ResourceManagerConfig ResourceManager { get; } = new ResourceManagerConfig();
     }
 }
diff --git a/src/Microsoft.Health.Fhir.Core/Configs/ResourceManagerConfig.cs b/src/Microsoft.Health.Fhir.Core/Configs/ResourceManagerConfig.cs
@@ -0,0 +1,11 @@
+// -------------------------------------------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+// -------------------------------------------------------------------------------------------------
+
+namespace Microsoft.Health.Fhir.Core.Configs;
+
+public class ResourceManagerConfig
+{
+    public string DataStoreResourceId { get; set; }
+}
diff --git a/src/Microsoft.Health.Fhir.CosmosDb.Core/Configs/CosmosDataStoreConfiguration.cs b/src/Microsoft.Health.Fhir.CosmosDb.Core/Configs/CosmosDataStoreConfiguration.cs
@@ -14,6 +14,16 @@ public class CosmosDataStoreConfiguration
 
         public string Key { get; set; }
 
+        /// <summary>
+        /// When true, the application will use the managed identity of the Azure resource to authenticate to Cosmos DB.
+        /// Key configuration will be ignored if this is set to true.
+        /// </summary>
+        public bool UseManagedIdentity { get; set; }
+
+        public bool AllowDatabaseCreation { get; set; } = true;
+
+        public bool AllowCollectionSetup { get; set; } = true;
+
         public string DatabaseId { get; set; }
 
         public int? InitialDatabaseThroughput { get; set; }
@@ -22,8 +32,6 @@ public class CosmosDataStoreConfiguration
 
         public ConsistencyLevel? DefaultConsistencyLevel { get; set; }
 
-        public bool AllowDatabaseCreation { get; set; } = true;
-
         [System.Diagnostics.CodeAnalysis.SuppressMessage("Usage", "CA2227:Collection properties should be read only", Justification = "This is a configuration class")]
         public IList<string> PreferredLocations { get; set; }
 

diff --git a/src/Microsoft.Health.Fhir.CosmosDb.Core/Features/Storage/ICollectionSetup.cs b/src/Microsoft.Health.Fhir.CosmosDb.Core/Features/Storage/ICollectionSetup.cs
@@ -6,6 +6,7 @@
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage.Versioning;
 using Polly;
 
 namespace Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage
@@ -16,6 +17,8 @@ public interface ICollectionSetup
 
         public Task CreateCollectionAsync(IEnumerable<ICollectionInitializer> collectionInitializers, AsyncPolicy retryPolicy, CancellationToken cancellationToken = default);
 
-        public Task UpdateFhirCollectionSettingsAsync(CancellationToken cancellationToken);
+        Task InstallStoredProcs(CancellationToken cancellationToken);
+
+        public Task UpdateFhirCollectionSettingsAsync(CollectionVersion version, CancellationToken cancellationToken);
     }
 }
diff --git a/src/Microsoft.Health.Fhir.CosmosDb.Core/Features/Storage/ICosmosClientTestProvider.cs b/src/Microsoft.Health.Fhir.CosmosDb.Core/Features/Storage/ICosmosClientTestProvider.cs
@@ -12,6 +12,6 @@ namespace Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage
 {
     public interface ICosmosClientTestProvider
     {
-        Task PerformTestAsync(Container container, CosmosDataStoreConfiguration configuration, CosmosCollectionConfiguration cosmosCollectionConfiguration, CancellationToken cancellationToken = default);
+        Task PerformTestAsync(Container container, CancellationToken cancellationToken = default);
     }
 }
diff --git a/src/Microsoft.Health.Fhir.CosmosDb.Initialization/Features/Storage/CollectionInitializer.cs b/src/Microsoft.Health.Fhir.CosmosDb.Initialization/Features/Storage/CollectionInitializer.cs
@@ -53,15 +53,16 @@ public async Task<Container> InitializeCollectionAsync(CosmosClient client, Asyn
 
             _logger.LogInformation("Finding Container: {CollectionId}", _cosmosCollectionConfiguration.CollectionId);
 
-            var existingContainer = await retryPolicy.ExecuteAsync(async () => await database.TryGetContainerAsync(_cosmosCollectionConfiguration.CollectionId));
+            ContainerResponse existingContainer = await retryPolicy.ExecuteAsync(async () => await database.TryGetContainerAsync(_cosmosCollectionConfiguration.CollectionId));
             _logger.LogInformation("Creating Cosmos Container if not exits: {CollectionId}", _cosmosCollectionConfiguration.CollectionId);
 
             ContainerResponse containerResponse = await retryPolicy
                                                         .ExecuteAsync(async () =>
                                                         await database.CreateContainerIfNotExistsAsync(
                                                         _cosmosCollectionConfiguration.CollectionId,
                                                         $"/{KnownDocumentProperties.PartitionKey}",
-                                                        _cosmosCollectionConfiguration.InitialCollectionThroughput));
+                                                        _cosmosCollectionConfiguration.InitialCollectionThroughput,
+                                                        cancellationToken: CancellationToken.None));
 
             if (containerResponse.StatusCode == HttpStatusCode.Created || containerResponse.Resource.DefaultTimeToLive != -1)
             {
@@ -79,7 +80,7 @@ await retryPolicy.ExecuteAsync(async () =>
             {
                 try
                 {
-                    await _clientTestProvider.PerformTestAsync(existingContainer, _cosmosDataStoreConfiguration, _cosmosCollectionConfiguration, cancellationToken);
+                    await _clientTestProvider.PerformTestAsync(existingContainer, cancellationToken);
                 }
                 catch (CosmosException e) when (e.StatusCode == HttpStatusCode.TooManyRequests)
                 {

diff --git a/...icrosoft.Health.Fhir.CosmosDb.Initialization/Features/Storage/DataPlaneCollectionSetup.cs b/...icrosoft.Health.Fhir.CosmosDb.Initialization/Features/Storage/DataPlaneCollectionSetup.cs
@@ -5,29 +5,24 @@
 
 using System;
 using System.Collections.Generic;
-using System.Linq;
-using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using EnsureThat;
 using Microsoft.Azure.Cosmos;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using Microsoft.Health.Abstractions.Exceptions;
-using Microsoft.Health.Core.Features.Context;
 using Microsoft.Health.Fhir.CosmosDb.Core;
 using Microsoft.Health.Fhir.CosmosDb.Core.Configs;
 using Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage;
 using Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage.StoredProcedures;
 using Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage.Versioning;
 using Polly;
-using Polly.Retry;
 
 namespace Microsoft.Health.Fhir.CosmosDb.Initialization.Features.Storage
 {
     public class DataPlaneCollectionSetup : ICollectionSetup
     {
-        private const int CollectionSettingsVersion = 3;
         private readonly ILogger<DataPlaneCollectionSetup> _logger;
         private readonly CosmosClient _client;
         private readonly Lazy<Container> _container;
@@ -105,12 +100,6 @@ public async Task CreateCollectionAsync(IEnumerable<ICollectionInitializer> coll
                 }
 
                 _logger.LogInformation("Collections successfully initialized");
-
-                _logger.LogInformation("Installing Stored Procedures");
-
-                await _storedProcedureInstaller.ExecuteAsync(_container.Value, cancellationToken);
-
-                _logger.LogInformation("Stored Procedures are installed");
             }
             catch (Exception ex)
             {
@@ -120,13 +109,20 @@ public async Task CreateCollectionAsync(IEnumerable<ICollectionInitializer> coll
             }
         }
 
-        public async Task UpdateFhirCollectionSettingsAsync(CancellationToken cancellationToken)
+        public async Task InstallStoredProcs(CancellationToken cancellationToken)
         {
-            EnsureArg.IsNotNull(_container, nameof(_container));
+            _logger.LogInformation("Installing Stored Procedures");
+
+            await _storedProcedureInstaller.ExecuteAsync(_container.Value, cancellationToken);
 
-            var thisVersion = await GetLatestCollectionVersionAsync(_container.Value, cancellationToken);
+            _logger.LogInformation("Stored Procedures are installed");
+        }
+
+        public async Task UpdateFhirCollectionSettingsAsync(CollectionVersion version, CancellationToken cancellationToken)
+        {
+            EnsureArg.IsNotNull(_container, nameof(_container));
 
-            if (thisVersion.Version < 2)
+            if (version.Version < 2)
             {
                 var containerResponse = await _container.Value.ReadContainerAsync(cancellationToken: cancellationToken);
 
@@ -149,24 +145,7 @@ public async Task UpdateFhirCollectionSettingsAsync(CancellationToken cancellati
                 containerResponse.Resource.DefaultTimeToLive = -1;
 
                 await _container.Value.ReplaceContainerAsync(containerResponse, cancellationToken: cancellationToken);
-
-                thisVersion.Version = CollectionSettingsVersion;
-                await _container.Value.UpsertItemAsync(thisVersion, new PartitionKey(thisVersion.PartitionKey), cancellationToken: cancellationToken);
             }
         }
-
-        private static async Task<CollectionVersion> GetLatestCollectionVersionAsync(Container container, CancellationToken cancellationToken)
-        {
-            FeedIterator<CollectionVersion> query = container.GetItemQueryIterator<CollectionVersion>(
-                new QueryDefinition("SELECT * FROM root r"),
-                requestOptions: new QueryRequestOptions
-                {
-                    PartitionKey = new PartitionKey(CollectionVersion.CollectionVersionPartition),
-                });
-
-            FeedResponse<CollectionVersion> result = await query.ReadNextAsync(cancellationToken);
-
-            return result.FirstOrDefault() ?? new CollectionVersion();
-        }
     }
 }