move cosmosdb database creation back to bicep

Author: jgbradley1

backend/graphrag_app/main.py

@@ -47,21 +47,17 @@ async def catch_all_exceptions_middleware(request: Request, call_next):
 
 
 def intialize_cosmosdb_setup():
-    """Initialise CosmosDB (if necessary) by setting up a database and containers that are expected at startup time."""
+    """Initialise database setup (if necessary) and configure CosmosDB containers that are expected at startup time if they do not exist."""
     azure_client_manager = AzureClientManager()
     client = azure_client_manager.get_cosmos_client()
     db_client = client.create_database_if_not_exists("graphrag")
     # create containers with default settings
-    throughput = ThroughputProperties(
-        auto_scale_max_throughput=1000, auto_scale_increment_percent=1
-    )
     db_client.create_container_if_not_exists(
-        id="jobs", partition_key=PartitionKey(path="/id"), offer_throughput=throughput
+        id="jobs", partition_key=PartitionKey(path="/id")
     )
     db_client.create_container_if_not_exists(
         id="container-store",
         partition_key=PartitionKey(path="/id"),
-        offer_throughput=throughput,
     )
 
 

infra/core/cosmosdb/cosmosdb.bicep

@@ -10,6 +10,8 @@ param location string = resourceGroup().location
 @allowed(['Enabled', 'Disabled'])
 param publicNetworkAccess string = 'Disabled'
 
+var maxThroughput = 1000
+
 resource cosmosDb 'Microsoft.DocumentDB/databaseAccounts@2024-11-15' = {
   name: cosmosDbName
   location: location
@@ -64,7 +66,25 @@ resource cosmosDb 'Microsoft.DocumentDB/databaseAccounts@2024-11-15' = {
     }
     networkAclBypassResourceIds: []
     capacity: {
-      totalThroughputLimit: 4000
+      totalThroughputLimit: maxThroughput
+    }
+  }
+}
+
+// create a single database that is used to maintain state information for graphrag indexing
+// NOTE: The current CosmosDB role assignments are not sufficient to allow the aks workload identity to create databases so we must do it in bicep at deployment time.
+// TODO: Identify and assign appropriate RBAC roles that allow the workload identity to create new databases instead of relying on this bicep implementation.
+resource graphragDatabase 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases@2024-11-15' = {
+  parent: cosmosDb
+  name: 'graphrag'
+  properties: {
+    options: {
+      autoscaleSettings: {
+        maxThroughput: maxThroughput
+      }
+    }
+    resource: {
+      id: 'graphrag'
     }
   }
 }

infra/core/rbac/workload-identity-rbac.bicep

@@ -31,6 +31,9 @@ var roleDefinitions = [
   {
     id: '3913510d-42f4-4e42-8a64-420c390055eb' // Monitoring Metrics Publisher Role
   }
+  {
+    id: '5bd9cd88-fe45-4216-938b-f97437e15450' // DocumentDB Account Contributor - enables control plane operations
+  }
 ]
 
 resource roleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = [
@@ -50,10 +53,24 @@ resource cosmosDb 'Microsoft.DocumentDB/databaseAccounts@2024-12-01-preview' exi
   name: cosmosDbName
 }
 
-var customRoleName = 'Custom cosmosDB role for graphrag - adds read/write permissions at the database and container level'
+// NOTE: The code snippet below is commented out because there is a known race condition issue at deployment time when assigning Cosmos DB built-in roles to an identity.
+// For more information: https://github.com/pulumi/pulumi-azure-native/issues/2816
+// For a temporary workaround, that seems to work in practice, we can create a custom role defintion with the same permissions as the built-in role and use it instead
+// var cosmosDbContainerReadWriteRoleId = '00000000-0000-0000-0000-000000000002'
+// resource sqlRoleAssignment 'Microsoft.DocumentDB/databaseAccounts/sqlRoleAssignments@2023-11-15' = {
+//   name: guid(cosmosDb.id, principalId, principalType, cosmosDbContainerReadWriteRoleId)
+//   parent: cosmosDb
+//   properties: {
+//     principalId: principalId
+//     roleDefinitionId: '${cosmosDb.id}/sqlRoleDefinitions/${cosmosDbContainerReadWriteRoleId}'
+//     scope: cosmosDb.id
+//   }
+// }
+
+var customRoleName = 'Custom cosmosDB role for graphrag - adds read/write permissions at the container level'
 resource customCosmosRoleDefinition 'Microsoft.DocumentDB/databaseAccounts/sqlRoleDefinitions@2024-12-01-preview' = {
   // note: the guid must be globally unique and deterministic (reproducible) across Azure
-  name: guid(subscription().subscriptionId, resourceGroup().name, cosmosDb.id, customRoleName) // guid is used to ensure uniqueness
+  name: guid(cosmosDb.id, customRoleName)
   parent: cosmosDb
   properties: {
     roleName: customRoleName
@@ -67,7 +84,6 @@ resource customCosmosRoleDefinition 'Microsoft.DocumentDB/databaseAccounts/sqlRo
           'Microsoft.DocumentDB/databaseAccounts/readMetadata'
           'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/*'
           'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/items/*'
-          'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/write'
         ]
       }
     ]
@@ -76,13 +92,7 @@ resource customCosmosRoleDefinition 'Microsoft.DocumentDB/databaseAccounts/sqlRo
 
 resource assignment 'Microsoft.DocumentDB/databaseAccounts/sqlRoleAssignments@2024-12-01-preview' = {
   // note: the guid must be globally unique and deterministic (reproducible) across Azure
-  name: guid(
-    subscription().subscriptionId,
-    resourceGroup().name,
-    cosmosDb.id,
-    customCosmosRoleDefinition.id,
-    principalId
-  )
+  name: guid(cosmosDb.id, principalId, principalType, customCosmosRoleDefinition.id)
   parent: cosmosDb
   properties: {
     principalId: principalId