Move cosmosdb database creation to bicep deployment and fix AI Search private endpoint (#255)

2025-06-27 04:39:57 +00:00 · 2025-03-10 19:41:28 -04:00 · 2025-03-10 19:41:28 -04:00 · 942b2c63a1
commit 942b2c63a1
parent bd11643a9f
22 changed files with 355 additions and 158 deletions
--- a/backend/graphrag_app/main.py
+++ b/backend/graphrag_app/main.py
@ -46,22 +46,24 @@ async def catch_all_exceptions_middleware(request: Request, call_next):
        return Response("Unexpected internal server error.", status_code=500)


+# NOTE: this function is not currently used, but it is a placeholder for future use once RBAC issues have been resolved
 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 = client.create_database_if_not_exists(
+        "graphrag", offer_throughput=throughput
+    )
+    # create containers with default settings
    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,
    )


@ -78,8 +80,8 @@ async def lifespan(app: FastAPI):
        yield
        return

-    # Initialize CosmosDB setup
-    intialize_cosmosdb_setup()
+    # TODO: must identify proper CosmosDB RBAC roles before databases and containers can be created by this web app
+    # intialize_cosmosdb_setup()

    try:
        # Check if the cronjob exists and create it if it does not exist
--- a/infra/abbreviations.json
+++ b/infra/abbreviations.json
@ -84,6 +84,7 @@
  "networkNetworkSecurityGroupsSecurityRules": "nsgsr-",
  "networkNetworkWatchers": "nw-",
  "networkPrivateDnsZones": "pdnsz-",
+  "networkPrivateLinkScope": "pls-",
  "networkPrivateLinkServices": "pl-",
  "networkPublicIPAddresses": "pip-",
  "networkPublicIPPrefixes": "ippre-",
--- a/infra/core/ai-search/ai-search.bicep
+++ b/infra/core/ai-search/ai-search.bicep
@ -10,7 +10,7 @@ param location string = resourceGroup().location
@allowed(['enabled', 'disabled'])
 param publicNetworkAccess string = 'enabled'

-resource aiSearch 'Microsoft.Search/searchServices@2024-03-01-preview' = {
+resource search 'Microsoft.Search/searchServices@2024-06-01-preview' = {
  name: name
  location: location
  sku: {
@ -21,9 +21,13 @@ resource aiSearch 'Microsoft.Search/searchServices@2024-03-01-preview' = {
    replicaCount: 1
    partitionCount: 1
    publicNetworkAccess: publicNetworkAccess
+    networkRuleSet: {
+      ipRules: []
+      bypass: 'AzureServices'
+    }
    semanticSearch: 'disabled'
  }
 }

-output name string = aiSearch.name
-output id string = aiSearch.id
+output name string = search.name
+output id string = search.id
--- a/infra/core/aks/aks.bicep
+++ b/infra/core/aks/aks.bicep
@ -50,7 +50,7 @@ param subnetId string

 param privateDnsZoneName string

-@description('Array of object ids that will have admin role of the cluster')
+@description('Array of object ids of admins that will have admin control over the cluster')
 param clusterAdmins array = []

 resource privateDnsZone 'Microsoft.Network/privateDnsZones@2020-06-01' existing = {
@ -187,11 +187,10 @@ resource aksManagedAutoUpgradeSchedule 'Microsoft.ContainerService/managedCluste
      schedule: {
        weekly: {
          intervalWeeks: 1
-          dayOfWeek: 'Monday'
+          dayOfWeek: 'Sunday'
        }
      }
      durationHours: 4
-      startDate: '2024-06-11'
      startTime: '12:00'
    }
  }
@ -209,7 +208,6 @@ resource aksManagedNodeOSUpgradeSchedule 'Microsoft.ContainerService/managedClus
        }
      }
      durationHours: 4
-      startDate: '2024-06-11'
      startTime: '12:00'
    }
  }
--- a/infra/core/apim/apim.bicep
+++ b/infra/core/apim/apim.bicep
@ -2,7 +2,7 @@
 // Licensed under the MIT License.

@description('The name of the API Management service instance')
-param apiManagementName string = 'apiservice${uniqueString(resourceGroup().id)}'
+param apiManagementName string

@description('The email address of the owner of the service')
@minLength(1)
--- a/infra/core/cosmosdb/cosmosdb.bicep
+++ b/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,101 @@ 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 and containers 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 and containers 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'
+    }
+  }
+}
+
+resource jobsContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers@2024-11-15' = {
+  parent: graphragDatabase
+  name: 'jobs'
+  properties: {
+    resource: {
+      id: 'jobs'
+      indexingPolicy: {
+        indexingMode: 'consistent'
+        automatic: true
+        includedPaths: [
+          {
+            path: '/*'
+          }
+        ]
+        excludedPaths: [
+          {
+            path: '/"_etag"/?'
+          }
+        ]
+      }
+      partitionKey: {
+        paths: [
+          '/id'
+        ]
+        kind: 'Hash'
+        version: 2
+      }
+      uniqueKeyPolicy: {
+        uniqueKeys: []
+      }
+      conflictResolutionPolicy: {
+        mode: 'LastWriterWins'
+        conflictResolutionPath: '/_ts'
+      }
+    }
+  }
+}
+
+resource containerStoreContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers@2024-11-15' = {
+  parent: graphragDatabase
+  name: 'container-store'
+  properties: {
+    resource: {
+      id: 'container-store'
+      indexingPolicy: {
+        indexingMode: 'consistent'
+        automatic: true
+        includedPaths: [
+          {
+            path: '/*'
+          }
+        ]
+        excludedPaths: [
+          {
+            path: '/"_etag"/?'
+          }
+        ]
+      }
+      partitionKey: {
+        paths: [
+          '/id'
+        ]
+        kind: 'Hash'
+        version: 2
+      }
+      uniqueKeyPolicy: {
+        uniqueKeys: []
+      }
+      conflictResolutionPolicy: {
+        mode: 'LastWriterWins'
+        conflictResolutionPath: '/_ts'
+      }
    }
  }
 }
--- a/infra/core/identity/identity.bicep
+++ b/infra/core/identity/identity.bicep
@ -24,5 +24,6 @@ resource federatedCredentialResources 'Microsoft.ManagedIdentity/userAssignedIde
 ]

 output name string = identity.name
+output id string = identity.id
 output clientId string = identity.properties.clientId
 output principalId string = identity.properties.principalId
--- a/infra/core/monitor/app-insights.bicep
+++ b/infra/core/monitor/app-insights.bicep
@ -25,6 +25,7 @@ resource appInsights 'Microsoft.Insights/components@2020-02-02' = {
  }
 }

+output name string = appInsights.name
 output id string = appInsights.id
 output connectionString string = appInsights.properties.ConnectionString
 output instrumentationKey string = appInsights.properties.InstrumentationKey
--- a/infra/core/monitor/private-link-scope.bicep
+++ b/infra/core/monitor/private-link-scope.bicep
@ -6,7 +6,7 @@ param privateLinkScopedResources array = []
 param queryAccessMode string = 'Open'
 param ingestionAccessMode string = 'PrivateOnly'

-resource privateLinkScope 'microsoft.insights/privateLinkScopes@2021-07-01-preview' = {
+resource privateLinkScope 'microsoft.Insights/privateLinkScopes@2021-07-01-preview' = {
  name: privateLinkScopeName
  location: 'global'
  properties: {
@ -17,7 +17,7 @@ resource privateLinkScope 'microsoft.insights/privateLinkScopes@2021-07-01-previ
  }
 }

-resource scopedResources 'microsoft.insights/privateLinkScopes/scopedResources@2021-07-01-preview' = [
+resource scopedResources 'Microsoft.Insights/privateLinkScopes/scopedResources@2021-07-01-preview' = [
  for id in privateLinkScopedResources: {
    name: uniqueString(id)
    parent: privateLinkScope
--- a/infra/core/rbac/workload-identity-rbac.bicep
+++ b/infra/core/rbac/workload-identity-rbac.bicep
@ -8,85 +8,189 @@ param principalId string
@allowed(['ServicePrincipal', 'User', 'Group', 'Device', 'ForeignGroup'])
 param principalType string

+@description('Name of an existing AI Search resource.')
+param aiSearchName string
+
+@description('Name of an existing AppInsights resource.')
+param appInsightsName string
+
@description('Name of an existing CosmosDB resource.')
 param cosmosDbName string

-@description('Role definitions for various roles that will be assigned at deployment time. Learn more: https://learn.microsoft.com/en-us/azure/role-based-access-control/built-in-roles')
-var roleDefinitions = [
-  {
-    id: 'ba92f5b4-2d11-453d-a403-e96b0029c9fe' // Storage Blob Data Contributor Role
-  }
-  {
-    id: 'b24988ac-6180-42a0-ab88-20f7382dd24c' // AI Search Contributor Role
-  }
-  {
-    id: '8ebe5a00-799e-43f5-93ac-243d3dce84a7' // AI Search Index Data Contributor Role
-  }
-  {
-    id: '1407120a-92aa-4202-b7e9-c0e197c71c8f' // AI Search Index Data Reader Role
-  }
-  {
-    id: 'a001fd3d-188f-4b5d-821b-7da978bf7442' // Cognitive Services OpenAI Contributor
-  }
-  {
-    id: '3913510d-42f4-4e42-8a64-420c390055eb' // Monitoring Metrics Publisher Role
-  }
-]
+@description('Name of an existing Azure Storage resource.')
+param storageName string

-resource roleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = [
-  for roleDef in roleDefinitions: {
-    // note: the guid must be globally unique and deterministic (reproducible) across Azure
-    name: guid(subscription().subscriptionId, resourceGroup().name, principalId, principalType, roleDef.id)
-    scope: resourceGroup()
-    properties: {
-      principalId: principalId
-      principalType: principalType
-      roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleDef.id)
-    }
-  }
-]
+@description('Role definitions for various roles that will be assigned. Learn more: https://learn.microsoft.com/en-us/azure/role-based-access-control/built-in-roles')
+var roleIds = {
+  contributor: 'b24988ac-6180-42a0-ab88-20f7382dd24c' // Contributor Role
+  aiSearchIndexDataContributor: '8ebe5a00-799e-43f5-93ac-243d3dce84a7' // AI Search Index Data Contributor Role
+  aiSearchIndexDataReader: '1407120a-92aa-4202-b7e9-c0e197c71c8f' // AI Search Index Data Reader Role
+  cognitiveServicesOpenAIContributor: 'a001fd3d-188f-4b5d-821b-7da978bf7442' // Cognitive Services OpenAI Contributor Role
+  cosmosDBOperator: '230815da-be43-4aae-9cb4-875f7bd000aa' // Cosmos DB Operator Role - cosmos control plane operations
+  cosmosDbBuiltInDataContributor: '00000000-0000-0000-0000-000000000002' // Cosmos Built-in Data Contributor Role - cosmos data plane operations
+  documentDBAccountContributor: '5bd9cd88-fe45-4216-938b-f97437e15450' // DocumentDB Account Contributor Role - cosmos control plane operations
+  monitoringMetricsPublisher: '3913510d-42f4-4e42-8a64-420c390055eb' // Monitoring Metrics Publisher Role
+  storageBlobDataContributor: 'ba92f5b4-2d11-453d-a403-e96b0029c9fe' // Storage Blob Data Contributor Role
+  sqlDBContributor: '9b7fa17d-e63e-47b0-bb0a-15c516ac86ec' // SQL DB Contributor Role - cosmos control plane operations
+}

-resource cosmosDb 'Microsoft.DocumentDB/databaseAccounts@2024-12-01-preview' existing = {
+// get references to existing resources
+resource aiSearch 'Microsoft.Search/searchServices@2024-03-01-preview' existing = {
+  name: aiSearchName
+}
+resource appInsights 'Microsoft.Insights/components@2020-02-02' existing = {
+  name: appInsightsName
+}
+resource cosmosDb 'Microsoft.DocumentDB/databaseAccounts@2024-11-15' existing = {
  name: cosmosDbName
 }
+resource storage 'Microsoft.Storage/storageAccounts@2023-01-01' existing = {
+  name: storageName
+}

-var customRoleName = 'Custom cosmosDB role for graphrag - adds read/write permissions at the database and 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
-  parent: cosmosDb
+// make RBAC role assignments to each resource
+resource contributorRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(aiSearch.id, principalId, principalType, roleIds.contributor)
+  scope: aiSearch
  properties: {
-    roleName: customRoleName
-    type: 'CustomRole'
-    assignableScopes: [
-      cosmosDb.id
-    ]
-    permissions: [
-      {
-        dataActions: [
-          'Microsoft.DocumentDB/databaseAccounts/readMetadata'
-          'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/*'
-          'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/items/*'
-          'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/write'
-        ]
-      }
-    ]
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.contributor)
  }
 }

-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
-  )
+resource cognitiveServicesOpenAIContributorRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(resourceGroup().id, principalId, principalType, roleIds.cognitiveServicesOpenAIContributor)
+  scope: resourceGroup()
+  properties: {
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.cognitiveServicesOpenAIContributor)
+  }
+}
+
+resource aiSearchIndexDataContributorRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(aiSearch.id, principalId, principalType, roleIds.aiSearchIndexDataContributor)
+  scope: aiSearch
+  properties: {
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.aiSearchIndexDataContributor)
+  }
+}
+
+resource aiSearchIndexDataReaderRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(aiSearch.id, principalId, principalType, roleIds.aiSearchIndexDataReader)
+  scope: aiSearch
+  properties: {
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.aiSearchIndexDataReader)
+  }
+}
+
+resource cosmosDbOperatorRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(cosmosDb.id, principalId, principalType, roleIds.cosmosDBOperator)
+  scope: cosmosDb
+  properties: {
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.cosmosDBOperator)
+  }
+}
+
+resource documentDbAccountContributorRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(cosmosDb.id, principalId, principalType, roleIds.documentDBAccountContributor)
+  scope: cosmosDb
+  properties: {
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.documentDBAccountContributor)
+  }
+}
+
+resource sqlDbContributorRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(cosmosDb.id, principalId, principalType, roleIds.sqlDBContributor)
+  scope: cosmosDb
+  properties: {
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.sqlDBContributor)
+  }
+}
+
+resource storageBlobDataContributorRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(storage.id, principalId, principalType, roleIds.storageBlobDataContributor)
+  scope: storage
+  properties: {
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.storageBlobDataContributor)
+  }
+}
+
+resource monitoringMetricsPublisherRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
+  // note: the guid must be globally unique and deterministic across Azure
+  name: guid(appInsights.id, principalId, roleIds.monitoringMetricsPublisher)
+  scope: appInsights
+  properties: {
+    principalId: principalId
+    principalType: principalType
+    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleIds.monitoringMetricsPublisher)
+  }
+}
+
+// NOTE: The SQL role assignment below can be flaky due to 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
+// In practice, one option that may not have such flaky behavior is to create a custom role defintion with the same permissions as the built-in role and use it instead
+resource sqlRoleAssignment 'Microsoft.DocumentDB/databaseAccounts/sqlRoleAssignments@2024-11-15' = {
+  name: guid(cosmosDb.id, principalId, principalType, roleIds.cosmosDbBuiltInDataContributor)
  parent: cosmosDb
  properties: {
    principalId: principalId
-    roleDefinitionId: customCosmosRoleDefinition.id
+    roleDefinitionId: '${cosmosDb.id}/sqlRoleDefinitions/${roleIds.cosmosDbBuiltInDataContributor}'
    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(cosmosDb.id, customRoleName)
+//   parent: cosmosDb
+//   properties: {
+//     roleName: customRoleName
+//     type: 'CustomRole'
+//     assignableScopes: [
+//       cosmosDb.id
+//     ]
+//     permissions: [
+//       {
+//         dataActions: [
+//           'Microsoft.DocumentDB/databaseAccounts/readMetadata'
+//           'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/*'
+//           'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/items/*'
+//         ]
+//       }
+//     ]
+//   }
+// }
+
+// resource customRoleAssignment 'Microsoft.DocumentDB/databaseAccounts/sqlRoleAssignments@2024-12-01-preview' = {
+//   // note: the guid must be globally unique and deterministic (reproducible) across Azure
+//   name: guid(cosmosDb.id, principalId, principalType, customCosmosRoleDefinition.id)
+//   parent: cosmosDb
+//   properties: {
+//     principalId: principalId
+//     roleDefinitionId: customCosmosRoleDefinition.id
+//     scope: cosmosDb.id
+//   }
+// }
--- a/infra/core/vnet/nsg.bicep
+++ b/infra/core/vnet/nsg.bicep
@ -7,7 +7,7 @@ param nsgName string = 'apim-nsg-${uniqueString(resourceGroup().id)}'
@description('Azure region where the resources will be deployed')
 param location string = resourceGroup().location

-resource nsg 'Microsoft.Network/networkSecurityGroups@2024-01-01' = {
+resource nsg 'Microsoft.Network/networkSecurityGroups@2024-05-01' = {
  name: nsgName
  location: location
  properties: {
--- a/infra/core/vnet/private-dns-vnet-link.bicep
+++ b/infra/core/vnet/private-dns-vnet-link.bicep
@ -5,13 +5,13 @@ param vnetId string
 param privateDnsZoneName string
 var vnet_id_hash = uniqueString(vnetId)

-resource dnsZone 'Microsoft.Network/privateDnsZones@2020-06-01' = {
+resource dnsZone 'Microsoft.Network/privateDnsZones@2024-06-01' = {
  name: privateDnsZoneName
  location: 'global'
  properties: {}
 }

-resource dnsZoneLinks 'Microsoft.Network/privateDnsZones/virtualNetworkLinks@2020-06-01' = {
+resource dnsZoneLinks 'Microsoft.Network/privateDnsZones/virtualNetworkLinks@2024-06-01' = {
  name: 'vnet-link-${privateDnsZoneName}-${vnet_id_hash}'
  location: 'global'
  parent: dnsZone
--- a/infra/core/vnet/private-dns-zone-a-record.bicep
+++ b/infra/core/vnet/private-dns-zone-a-record.bicep
@ -13,11 +13,11 @@ param ttl int = 900
@description('The IP address')
 param ipv4Address string

-resource dnsZone 'Microsoft.Network/privateDnsZones@2020-06-01' existing = {
+resource dnsZone 'Microsoft.Network/privateDnsZones@2024-06-01' existing = {
  name: dnsZoneName
 }

-resource aRecord 'Microsoft.Network/privateDnsZones/A@2020-06-01' = {
+resource aRecord 'Microsoft.Network/privateDnsZones/A@2024-06-01' = {
  name: name
  parent: dnsZone
  properties: {
--- a/infra/core/vnet/private-dns-zone-groups.json
+++ b/infra/core/vnet/private-dns-zone-groups.json
@ -1,12 +1,13 @@
 {
  "azureCloud": {
-    "aiSearch": "privatelink.search.azure.com",
+    "aiSearch": "privatelink.search.windows.net",
    "azureMonitor": [
      "privatelink.monitor.azure.com",
      "privatelink.oms.opinsights.azure.com",
      "privatelink.agentsvc.azure-automation.net",
      "privatelink.ods.opinsights.azure.com"
    ],
+    "blobStorage": "privatelink.blob.core.windows.net",
    "cosmosDB": "privatelink.documents.azure.com"
  },
  "azureusgovernment": {
@ -17,6 +18,7 @@
      "privatelink.agentsvc.azure-automation.us",
      "privatelink.ods.opinsights.azure.us"
    ],
+    "blobStorage": "privatelink.blob.core.usgovcloudapi.net",
    "cosmosDB": "privatelink.documents.azure.us"
  }
 }
--- a/infra/core/vnet/private-dns-zone.bicep
+++ b/infra/core/vnet/private-dns-zone.bicep
@ -5,33 +5,29 @@
 param name string

@description('The name of the virtual networks the DNS zone should be associated with.')
-param vnetNames string[]
+param vnetName string

-resource dnsZone 'Microsoft.Network/privateDnsZones@2020-06-01' = {
+resource vnet 'Microsoft.Network/virtualNetworks@2024-05-01' existing = {
+  name: vnetName
+}
+
+resource dnsZone 'Microsoft.Network/privateDnsZones@2024-06-01' = {
  name: name
  location: 'global'
  properties: {}
 }

-resource vnets 'Microsoft.Network/virtualNetworks@2024-01-01' existing = [
-  for vnetName in vnetNames: {
-    name: vnetName
-  }
-]
-
-resource dnsZoneLinks 'Microsoft.Network/privateDnsZones/virtualNetworkLinks@2020-06-01' = [
-  for (vnetName, index) in vnetNames: {
-    name: vnetName
-    location: 'global'
-    parent: dnsZone
-    properties: {
-      registrationEnabled: false
-      virtualNetwork: {
-        id: vnets[index].id
-      }
+resource dnsZoneLinks 'Microsoft.Network/privateDnsZones/virtualNetworkLinks@2024-06-01' = {
+  name: vnetName
+  location: 'global'
+  parent: dnsZone
+  properties: {
+    registrationEnabled: false
+    virtualNetwork: {
+      id: vnet.id
    }
  }
-]
+}

 output name string = dnsZone.name
 output id string = dnsZone.id
--- a/infra/core/vnet/private-endpoint.bicep
+++ b/infra/core/vnet/private-endpoint.bicep
@ -14,7 +14,7 @@ param privateEndpointName string
 param groupId string
 param location string = resourceGroup().location

-resource privateEndpoint 'Microsoft.Network/privateEndpoints@2021-05-01' = {
+resource privateEndpoint 'Microsoft.Network/privateEndpoints@2024-05-01' = {
  name: privateEndpointName
  location: location
  properties: {
@ -33,7 +33,7 @@ resource privateEndpoint 'Microsoft.Network/privateEndpoints@2021-05-01' = {
  }
 }

-resource privateDnsZoneGroup 'Microsoft.Network/privateEndpoints/privateDnsZoneGroups@2021-05-01' = {
+resource privateDnsZoneGroup 'Microsoft.Network/privateEndpoints/privateDnsZoneGroups@2024-05-01' = {
  name: groupId
  parent: privateEndpoint
  properties: {
--- a/infra/core/vnet/privatelink-private-dns-zones.bicep
+++ b/infra/core/vnet/privatelink-private-dns-zones.bicep
@ -1,14 +1,14 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.

-@description('Virtual Network IDs to link to')
-param linkedVnetIds array
+@description('The virtual network ID to link to')
+param linkedVnetId string

-var privateDnsZoneData = loadJsonContent('private-dns-zone-groups.json')
+var privateDnsZoneData = loadJsonContent('private-dns-zone-groups.json') // for more information: https://learn.microsoft.com/en-us/azure/azure-government/compare-azure-government-global-azure
 var cloudName = toLower(environment().name)

 var aiSearchPrivateDnsZoneName = privateDnsZoneData[cloudName].aiSearch
-var blobStoragePrivateDnsZoneName = 'privatelink.blob.${environment().suffixes.storage}'
+var blobStoragePrivateDnsZoneName = privateDnsZoneData[cloudName].blobStorage
 var cosmosDbPrivateDnsZoneName = privateDnsZoneData[cloudName].cosmosDb
 var storagePrivateDnsZoneNames = [blobStoragePrivateDnsZoneName]
 var azureMonitorPrivateDnsZones = privateDnsZoneData[cloudName].azureMonitor
@ -20,7 +20,7 @@ var privateDnsZones = union(
  [aiSearchPrivateDnsZoneName]
 )

-resource privateDnsZoneResources 'Microsoft.Network/privateDnsZones@2020-06-01' = [
+resource privateDnsZoneResources 'Microsoft.Network/privateDnsZones@2024-06-01' = [
  for name in privateDnsZones: {
    name: name
    location: 'global'
@ -32,7 +32,7 @@ module dnsVnetLinks 'vnet-dns-link.bicep' = [
    name: replace(privateDnsZoneName, '.', '-')
    params: {
      privateDnsZoneName: privateDnsZoneResources[index].name
-      vnetIds: linkedVnetIds
+      vnetId: linkedVnetId
    }
  }
 ]
--- a/infra/core/vnet/vnet-dns-link.bicep
+++ b/infra/core/vnet/vnet-dns-link.bicep
@ -2,22 +2,21 @@
 // Licensed under the MIT License.

 param privateDnsZoneName string
-param vnetIds array
+param vnetId string

-resource privateDnsZone 'Microsoft.Network/privateDnsZones@2020-06-01' existing = {
+resource privateDnsZone 'Microsoft.Network/privateDnsZones@2024-06-01' existing = {
  name: privateDnsZoneName
 }

-resource dnsVnetLinks 'Microsoft.Network/privateDnsZones/virtualNetworkLinks@2020-06-01' = [
-  for vnetId in vnetIds: {
-    name: '${replace(privateDnsZoneName, '.', '-')}-${uniqueString(vnetId)}'
-    parent: privateDnsZone
-    location: 'global'
-    properties: {
-      virtualNetwork: {
-        id: vnetId
-      }
-      registrationEnabled: false
+resource dnsVnetLinks 'Microsoft.Network/privateDnsZones/virtualNetworkLinks@2024-06-01' = {
+  name: '${replace(privateDnsZoneName, '.', '-')}-${uniqueString(vnetId)}'
+  parent: privateDnsZone
+  location: 'global'
+  properties: {
+    registrationEnabled: false
+    resolutionPolicy: 'Default'
+    virtualNetwork: {
+      id: vnetId
    }
  }
-]
+}
--- a/infra/core/vnet/vnet.bicep
+++ b/infra/core/vnet/vnet.bicep
@ -17,7 +17,7 @@ param apimTier string
@description('NSG resource ID.')
 param nsgID string

-resource vnet 'Microsoft.Network/virtualNetworks@2024-01-01' = {
+resource vnet 'Microsoft.Network/virtualNetworks@2024-05-01' = {
  name: vnetName
  location: location
  properties: {
@ -67,7 +67,7 @@ resource vnet 'Microsoft.Network/virtualNetworks@2024-01-01' = {
  }
 }

-output vnetId string = vnet.id
-output vnetName string = vnet.name
+output name string = vnet.name
+output id string = vnet.id
 output apimSubnetId string = vnet.properties.subnets[0].id
 output aksSubnetId string = vnet.properties.subnets[1].id
--- a/infra/deploy.sh
+++ b/infra/deploy.sh
@ -313,9 +313,6 @@ checkForApimSoftDelete () {

 deployAzureResources () {
    echo "Deploying Azure resources..."
-    # get principal/object id of the signed in user
-    local deployerPrincipalId=$(az ad signed-in-user show --output json | jq -r .id)
-    exitIfValueEmpty $deployerPrincipalId "Principal ID of deployer not found"
    local datetime="`date +%Y-%m-%d_%H-%M-%S`"
    local deployName="graphrag-deploy-$datetime"
    echo "Deployment name: $deployName"
@ -331,7 +328,6 @@ deployAzureResources () {
        --parameters "apiPublisherEmail=$PUBLISHER_EMAIL" \
        --parameters "enablePrivateEndpoints=$ENABLE_PRIVATE_ENDPOINTS" \
        --parameters "acrName=$CONTAINER_REGISTRY_NAME" \
-        --parameters "deployerPrincipalId=$deployerPrincipalId" \
        --output json)
    # errors in deployment may not be caught by exitIfCommandFailed function so we also check the output for errors
    exitIfCommandFailed $? "Error deploying Azure resources..."
--- a/infra/main.bicep
+++ b/infra/main.bicep
@ -31,9 +31,6 @@ var resourceBaseNameFinal = !empty(resourceBaseName)
@description('Cloud region for all resources')
 param location string = az.resourceGroup().location

-@description('Principal/Object ID of the deployer. Will be used to assign admin roles to the AKS cluster.')
-param deployerPrincipalId string
-
@minLength(1)
@description('Name of the publisher of the API Management instance.')
 param apiPublisherName string = 'Microsoft'
@ -45,7 +42,7 @@ param apiPublisherEmail string = 'publisher@microsoft.com'
@description('The AKS namespace to install GraphRAG in.')
 param aksNamespace string = 'graphrag'

-@description('Whether to enable private endpoints.')
+@description('Whether to use private endpoint connections or not.')
 param enablePrivateEndpoints bool = true

@description('Whether to restore the API Management instance.')
@ -92,7 +89,10 @@ module aksWorkloadIdentityRBAC 'core/rbac/workload-identity-rbac.bicep' = {
  params: {
    principalId: workloadIdentity.outputs.principalId
    principalType: 'ServicePrincipal'
+    aiSearchName: aiSearch.outputs.name
+    appInsightsName: appInsights.outputs.name
    cosmosDbName: cosmosdb.outputs.name
+    storageName: storage.outputs.name
  }
 }

@ -163,7 +163,7 @@ module aks 'core/aks/aks.bicep' = {
    location: location
    graphragVMSize: 'standard_d8s_v5' // 8 vcpu, 32 GB memory
    graphragIndexingVMSize: 'standard_e8s_v5' // 8 vcpus, 64 GB memory
-    clusterAdmins: !empty(deployerPrincipalId) ? ['${deployerPrincipalId}'] : null
+    clusterAdmins: [deployer().objectId]
    logAnalyticsWorkspaceId: log.outputs.id
    subnetId: vnet.outputs.aksSubnetId
    privateDnsZoneName: privateDnsZone.outputs.name
@ -260,25 +260,21 @@ module privateDnsZone 'core/vnet/private-dns-zone.bicep' = {
  name: 'private-dns-zone-deployment'
  params: {
    name: dnsDomain
-    vnetNames: [
-      vnet.outputs.vnetName // name
-    ]
+    vnetName: vnet.outputs.name
  }
 }

 module privatelinkPrivateDns 'core/vnet/privatelink-private-dns-zones.bicep' = if (enablePrivateEndpoints) {
  name: 'privatelink-private-dns-zones-deployment'
  params: {
-    linkedVnetIds: [
-      vnet.outputs.vnetId // id
-    ]
+    linkedVnetId: vnet.outputs.id
  }
 }

 module azureMonitorPrivateLinkScope 'core/monitor/private-link-scope.bicep' = if (enablePrivateEndpoints) {
  name: 'azure-monitor-privatelink-scope-deployment'
  params: {
-    privateLinkScopeName: 'pls-${resourceBaseNameFinal}'
+    privateLinkScopeName: '${abbrs.networkPrivateLinkScope}${resourceBaseNameFinal}'
    privateLinkScopedResources: [
      log.outputs.id
      appInsights.outputs.id
@ -334,6 +330,7 @@ module privateLinkScopePrivateEndpoint 'core/vnet/private-endpoint.bicep' = if (
  }
 }

+output deployer_principal_id string = deployer().objectId
 output azure_location string = location
 output azure_tenant_id string = tenant().tenantId
 output azure_ai_search_name string = aiSearch.outputs.name
--- a/notebooks/2-Advanced_Getting_Started.ipynb
+++ b/notebooks/2-Advanced_Getting_Started.ipynb
@ -61,7 +61,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
@ -100,7 +100,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
   "id": "7",
   "metadata": {
    "tags": []
@ -132,7 +132,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
   "id": "9",
   "metadata": {
    "tags": []
@ -159,7 +159,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
@ -181,7 +181,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
   "id": "12",
   "metadata": {
    "tags": []
@ -543,7 +543,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 35,
+   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
@ -566,7 +566,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
@ -702,7 +702,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 40,
+   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
@ -944,7 +944,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": null,
   "id": "54",
   "metadata": {},
   "outputs": [],