Skip to content
This repository was archived by the owner on May 27, 2025. It is now read-only.

Commit 97ea601

Browse files
committed
add custom cosmosdb rbac role
1 parent 1db953c commit 97ea601

File tree

9 files changed

+161
-230
lines changed

9 files changed

+161
-230
lines changed

backend/src/api/prompt_tuning.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ async def generate_prompts(storage_name: str, limit: int = 5):
4444
this_directory = os.path.dirname(
4545
os.path.abspath(inspect.getfile(inspect.currentframe()))
4646
)
47-
data = yaml.safe_load(open(f"{this_directory}/pipeline-settings.yaml"))
47+
data = yaml.safe_load(open(f"{this_directory}/../indexer/settings.yaml"))
4848
data["input"]["container_name"] = sanitized_storage_name
4949
graphrag_config = create_graphrag_config(values=data, root_dir=".")
5050

backend/src/main.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from contextlib import asynccontextmanager
77

88
import yaml
9+
from azure.cosmos import PartitionKey, ThroughputProperties
910
from fastapi import (
1011
FastAPI,
1112
Request,
@@ -49,10 +50,18 @@ def intialize_cosmosdb_setup():
4950
"""Initialise CosmosDB (if necessary) by setting up a database and containers that are expected at startup time."""
5051
azure_client_manager = AzureClientManager()
5152
client = azure_client_manager.get_cosmos_client()
52-
client.create_database_if_not_exists("graphrag")
53-
client.get_database_client("graphrag").create_container_if_not_exists("jobs", "/id")
54-
client.get_database_client("graphrag").create_container_if_not_exists(
55-
"container-store", "/id"
53+
db_client = client.create_database_if_not_exists("graphrag")
54+
# create containers with default settings
55+
throughput = ThroughputProperties(
56+
auto_scale_max_throughput=1000, auto_scale_increment_percent=1
57+
)
58+
db_client.create_container_if_not_exists(
59+
id="jobs", partition_key=PartitionKey(path="/id"), offer_throughput=throughput
60+
)
61+
db_client.create_container_if_not_exists(
62+
id="container-store",
63+
partition_key=PartitionKey(path="/id"),
64+
offer_throughput=throughput,
5665
)
5766

5867

infra/core/acr/acr.bicep

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
14
@description('The name of the Container Registry resource. Will be automatically generated if not provided.')
25
param registryName string
36

47
@description('The location of the Container Registry resource.')
58
param location string = resourceGroup().location
69

7-
@description('Array of objects with fields principalId, principalType, roleDefinitionId')
8-
param roleAssignments array = []
9-
1010
resource registry 'Microsoft.ContainerRegistry/registries@2023-11-01-preview' = {
1111
name: registryName
1212
location: location
@@ -27,14 +27,6 @@ resource registry 'Microsoft.ContainerRegistry/registries@2023-11-01-preview' =
2727
}
2828
}
2929

30-
resource roleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = [
31-
for role in roleAssignments: {
32-
name: guid('${role.principalId}-${role.principalType}-${role.roleDefinitionId}')
33-
scope: resourceGroup()
34-
properties: role
35-
}
36-
]
37-
3830
output name string = registry.name
3931
output id string = registry.id
4032
output loginServer string = registry.properties.loginServer

infra/core/aks/aks.bicep

Lines changed: 2 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -50,12 +50,6 @@ param subnetId string
5050

5151
param privateDnsZoneName string
5252

53-
@description('Array of objects with fields principalType, roleDefinitionId')
54-
param ingressRoleAssignments array = []
55-
56-
@description('Array of objects with fields principalType, roleDefinitionId')
57-
param systemRoleAssignments array = []
58-
5953
@description('Array of object ids that will have admin role of the cluster')
6054
param clusterAdmins array = []
6155

@@ -221,35 +215,11 @@ resource aksManagedNodeOSUpgradeSchedule 'Microsoft.ContainerService/managedClus
221215
}
222216
}
223217

224-
// role assignment to ingress identity
225-
resource webAppRoutingPrivateDnsContributor 'Microsoft.Authorization/roleAssignments@2022-04-01' = [
226-
for role in ingressRoleAssignments: {
227-
name: guid(subscription().subscriptionId, resourceGroup().name, role.roleDefinitionId, privateDnsZone.id)
228-
scope: resourceGroup()
229-
properties: {
230-
principalId: aks.properties.ingressProfile.webAppRouting.identity.objectId
231-
principalType: role.principalType
232-
roleDefinitionId: role.roleDefinitionId
233-
}
234-
}
235-
]
236-
237-
// role assignment to AKS system identity
238-
resource systemRoleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = [
239-
for role in systemRoleAssignments: {
240-
name: guid(subscription().subscriptionId, resourceGroup().name, role.roleDefinitionId, aks.id)
241-
scope: resourceGroup()
242-
properties: {
243-
principalId: aks.identity.principalId
244-
principalType: role.principalType
245-
roleDefinitionId: role.roleDefinitionId
246-
}
247-
}
248-
]
249-
250218
output name string = aks.name
251219
output id string = aks.id
252220
output managedResourceGroup string = aks.properties.nodeResourceGroup
253221
output controlPlaneFqdn string = aks.properties.fqdn
254222
output kubeletPrincipalId string = aks.properties.identityProfile.kubeletidentity.objectId
223+
output ingressWebAppIdentity string = aks.properties.ingressProfile.webAppRouting.identity.objectId
224+
output systemIdentity string = aks.identity.principalId
255225
output issuer string = aks.properties.oidcIssuerProfile.issuerURL

infra/core/cosmosdb/cosmosdb.bicep

Lines changed: 1 addition & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -7,22 +7,9 @@ param cosmosDbName string
77
@description('The location of the CosmosDB resource.')
88
param location string = resourceGroup().location
99

10-
@allowed([ 'Enabled', 'Disabled' ])
10+
@allowed(['Enabled', 'Disabled'])
1111
param publicNetworkAccess string = 'Disabled'
1212

13-
@description('Role definition id to assign to the principal. Learn more: https://learn.microsoft.com/en-us/azure/cosmos-db/how-to-setup-rbac')
14-
@allowed([
15-
'00000000-0000-0000-0000-000000000001' // 'Cosmos DB Built-in Data Reader' role
16-
'00000000-0000-0000-0000-000000000002' // 'Cosmos DB Built-in Data Contributor' role
17-
])
18-
param roleDefinitionId array = [
19-
'00000000-0000-0000-0000-000000000001'
20-
'00000000-0000-0000-0000-000000000002'
21-
]
22-
23-
param principalId string
24-
25-
2613
resource cosmosDb 'Microsoft.DocumentDB/databaseAccounts@2024-11-15' = {
2714
name: cosmosDbName
2815
location: location
@@ -82,104 +69,6 @@ resource cosmosDb 'Microsoft.DocumentDB/databaseAccounts@2024-11-15' = {
8269
}
8370
}
8471

85-
resource graphragDatabase 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases@2024-11-15' = {
86-
parent: cosmosDb
87-
name: 'graphrag'
88-
properties: {
89-
resource: {
90-
id: 'graphrag'
91-
}
92-
}
93-
}
94-
95-
resource jobsContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers@2024-11-15' = {
96-
parent: graphragDatabase
97-
name: 'jobs'
98-
properties: {
99-
resource: {
100-
id: 'jobs'
101-
indexingPolicy: {
102-
indexingMode: 'consistent'
103-
automatic: true
104-
includedPaths: [
105-
{
106-
path: '/*'
107-
}
108-
]
109-
excludedPaths: [
110-
{
111-
path: '/"_etag"/?'
112-
}
113-
]
114-
}
115-
partitionKey: {
116-
paths: [
117-
'/id'
118-
]
119-
kind: 'Hash'
120-
version: 2
121-
}
122-
uniqueKeyPolicy: {
123-
uniqueKeys: []
124-
}
125-
conflictResolutionPolicy: {
126-
mode: 'LastWriterWins'
127-
conflictResolutionPath: '/_ts'
128-
}
129-
}
130-
}
131-
}
132-
133-
resource containerStoreContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers@2024-11-15' = {
134-
parent: graphragDatabase
135-
name: 'container-store'
136-
properties: {
137-
resource: {
138-
id: 'container-store'
139-
indexingPolicy: {
140-
indexingMode: 'consistent'
141-
automatic: true
142-
includedPaths: [
143-
{
144-
path: '/*'
145-
}
146-
]
147-
excludedPaths: [
148-
{
149-
path: '/"_etag"/?'
150-
}
151-
]
152-
}
153-
partitionKey: {
154-
paths: [
155-
'/id'
156-
]
157-
kind: 'Hash'
158-
version: 2
159-
}
160-
uniqueKeyPolicy: {
161-
uniqueKeys: []
162-
}
163-
conflictResolutionPolicy: {
164-
mode: 'LastWriterWins'
165-
conflictResolutionPath: '/_ts'
166-
}
167-
}
168-
}
169-
}
170-
171-
resource sqlRoleAssignment 'Microsoft.DocumentDB/databaseAccounts/sqlRoleAssignments@2024-11-15' = [
172-
for roleId in roleDefinitionId: {
173-
name: guid('${roleId}-${principalId}-${cosmosDb.id}')
174-
parent: cosmosDb
175-
properties: {
176-
roleDefinitionId: '${resourceGroup().id}/providers/Microsoft.DocumentDB/databaseAccounts/${cosmosDb.name}/sqlRoleDefinitions/${roleId}'
177-
principalId: principalId
178-
scope: cosmosDb.id
179-
}
180-
}
181-
]
182-
18372
output name string = cosmosDb.name
18473
output id string = cosmosDb.id
18574
output endpoint string = cosmosDb.properties.documentEndpoint

infra/core/rbac/aks-rbac.bicep

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
4+
@description('Array of objects with fields principalId, principalType, roleDefinitionId')
5+
param roleAssignments array = []
6+
7+
resource roleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = [
8+
for role in roleAssignments: {
9+
// note: the guid must be globally unique and deterministic (reproducible) across Azure
10+
name: guid(
11+
subscription().subscriptionId,
12+
resourceGroup().name,
13+
role.principalId,
14+
role.principalType,
15+
role.roleDefinitionId
16+
)
17+
scope: resourceGroup()
18+
properties: role
19+
}
20+
]
Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
4+
@description('ID of the service principal to assign the RBAC roles to.')
5+
param principalId string
6+
7+
@description('Type of principal to assign the RBAC roles to.')
8+
@allowed(['ServicePrincipal', 'User', 'Group', 'Device', 'ForeignGroup'])
9+
param principalType string
10+
11+
@description('Name of an existing CosmosDB resource.')
12+
param cosmosDbName string
13+
14+
@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')
15+
var roleDefinitions = [
16+
{
17+
id: 'ba92f5b4-2d11-453d-a403-e96b0029c9fe' // Storage Blob Data Contributor Role
18+
}
19+
{
20+
id: 'b24988ac-6180-42a0-ab88-20f7382dd24c' // AI Search Contributor Role
21+
}
22+
{
23+
id: '8ebe5a00-799e-43f5-93ac-243d3dce84a7' // AI Search Index Data Contributor Role
24+
}
25+
{
26+
id: '1407120a-92aa-4202-b7e9-c0e197c71c8f' // AI Search Index Data Reader Role
27+
}
28+
{
29+
id: 'a001fd3d-188f-4b5d-821b-7da978bf7442' // Cognitive Services OpenAI Contributor
30+
}
31+
{
32+
id: '3913510d-42f4-4e42-8a64-420c390055eb' // Monitoring Metrics Publisher Role
33+
}
34+
]
35+
36+
resource roleAssignment 'Microsoft.Authorization/roleAssignments@2022-04-01' = [
37+
for roleDef in roleDefinitions: {
38+
// note: the guid must be globally unique and deterministic (reproducible) across Azure
39+
name: guid(subscription().subscriptionId, resourceGroup().name, principalId, principalType, roleDef.id)
40+
scope: resourceGroup()
41+
properties: {
42+
principalId: principalId
43+
principalType: principalType
44+
roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions', roleDef.id)
45+
}
46+
}
47+
]
48+
49+
resource cosmosDb 'Microsoft.DocumentDB/databaseAccounts@2024-12-01-preview' existing = {
50+
name: cosmosDbName
51+
}
52+
53+
var customRoleName = 'Custom cosmosDB role for graphrag - adds read/write permissions at the database and container level'
54+
resource customCosmosRoleDefinition 'Microsoft.DocumentDB/databaseAccounts/sqlRoleDefinitions@2024-12-01-preview' = {
55+
// note: the guid must be globally unique and deterministic (reproducible) across Azure
56+
name: guid(subscription().subscriptionId, resourceGroup().name, cosmosDb.id, customRoleName) // guid is used to ensure uniqueness
57+
parent: cosmosDb
58+
properties: {
59+
roleName: customRoleName
60+
type: 'CustomRole'
61+
assignableScopes: [
62+
cosmosDb.id
63+
]
64+
permissions: [
65+
{
66+
dataActions: [
67+
'Microsoft.DocumentDB/databaseAccounts/readMetadata'
68+
'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/*'
69+
'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/items/*'
70+
'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/write'
71+
]
72+
}
73+
]
74+
}
75+
}
76+
77+
resource assignment 'Microsoft.DocumentDB/databaseAccounts/sqlRoleAssignments@2024-12-01-preview' = {
78+
// note: the guid must be globally unique and deterministic (reproducible) across Azure
79+
name: guid(
80+
subscription().subscriptionId,
81+
resourceGroup().name,
82+
cosmosDb.id,
83+
customCosmosRoleDefinition.id,
84+
principalId
85+
)
86+
parent: cosmosDb
87+
properties: {
88+
principalId: principalId
89+
roleDefinitionId: customCosmosRoleDefinition.id
90+
scope: cosmosDb.id
91+
}
92+
}

0 commit comments

Comments
 (0)