Add schema support for cosmos db

cli/azd/internal/cmd/add/add_configure.go

@@ -44,6 +44,13 @@ func Configure(
 	case project.ResourceTypeDbPostgres,
 		project.ResourceTypeDbMongo:
 		return fillDatabaseName(ctx, r, console, p)
+	case project.ResourceTypeDbCosmos:
+		r, err := fillDatabaseName(ctx, r, console, p)
+		if err != nil {
+			return nil, err
+		}
+		r.Props = project.CosmosDBProps{}
+		return r, nil
 	case project.ResourceTypeMessagingEventHubs:
 		return fillEventHubs(ctx, r, console, p)
 	case project.ResourceTypeMessagingServiceBus:

cli/azd/internal/cmd/add/add_preview.go

@@ -67,6 +67,11 @@ func Metadata(r *project.ResourceConfig) resourceMeta {
 		res.UseEnvVars = []string{
 			"AZURE_OPENAI_ENDPOINT",
 		}
+	case project.ResourceTypeDbCosmos:
+		res.AzureResourceType = "Microsoft.DocumentDB/databaseAccounts"
+		res.UseEnvVars = []string{
+			"AZURE_COSMOS_ENDPOINT",
+		}
 	case project.ResourceTypeMessagingEventHubs:
 		res.AzureResourceType = "Microsoft.EventHub/namespaces"
 		res.UseEnvVars = []string{

cli/azd/internal/scaffold/scaffold_test.go

@@ -89,6 +89,9 @@ func TestExecInfra(t *testing.T) {
 				DbCosmosMongo: &DatabaseCosmosMongo{
 					DatabaseName: "appdb",
 				},
+				DbCosmos: &DatabaseCosmos{
+					DatabaseName: "cosmos",
+				},
 				DbRedis:        &DatabaseRedis{},
 				ServiceBus:     &ServiceBus{},
 				EventHubs:      &EventHubs{},
@@ -113,6 +116,9 @@ func TestExecInfra(t *testing.T) {
 						DbPostgres: &DatabaseReference{
 							DatabaseName: "appdb",
 						},
+						DbCosmos: &DatabaseReference{
+							DatabaseName: "cosmos",
+						},
 						ServiceBus:     &ServiceBus{},
 						EventHubs:      &EventHubs{},
 						StorageAccount: &StorageReference{},
@@ -181,6 +187,21 @@ func TestExecInfra(t *testing.T) {
 				},
 			},
 		},
+		{
+			"API with Cosmos",
+			InfraSpec{
+				DbCosmos: &DatabaseCosmos{},
+				Services: []ServiceSpec{
+					{
+						Name: "api",
+						Port: 3100,
+						DbCosmos: &DatabaseReference{
+							DatabaseName: "cosmos",
+						},
+					},
+				},
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

cli/azd/internal/scaffold/spec.go

@@ -15,6 +15,7 @@ type InfraSpec struct {
 	// Databases to create
 	DbPostgres    *DatabasePostgres
 	DbCosmosMongo *DatabaseCosmosMongo
+	DbCosmos      *DatabaseCosmos
 	DbRedis       *DatabaseRedis
 
 	// Messaging services
@@ -44,6 +45,16 @@ type DatabaseCosmosMongo struct {
 	DatabaseName string
 }
 
+type DatabaseCosmos struct {
+	DatabaseName string
+	Containers   []CosmosSqlDatabaseContainer
+}
+
+type CosmosSqlDatabaseContainer struct {
+	ContainerName     string
+	PartitionKeyPaths []string
+}
+
 type DatabaseRedis struct {
 }
 
@@ -89,6 +100,7 @@ type ServiceSpec struct {
 	// Connection to a database
 	DbPostgres    *DatabaseReference
 	DbCosmosMongo *DatabaseReference
+	DbCosmos      *DatabaseReference
 	DbRedis       *DatabaseReference
 
 	StorageAccount *StorageReference

cli/azd/pkg/project/resources.go

@@ -16,6 +16,7 @@ func AllResourceTypes() []ResourceType {
 		ResourceTypeDbRedis,
 		ResourceTypeDbPostgres,
 		ResourceTypeDbMongo,
+		ResourceTypeDbCosmos,
 		ResourceTypeHostContainerApp,
 		ResourceTypeOpenAiModel,
 		ResourceTypeMessagingEventHubs,
@@ -28,6 +29,7 @@ const (
 	ResourceTypeDbRedis             ResourceType = "db.redis"
 	ResourceTypeDbPostgres          ResourceType = "db.postgres"
 	ResourceTypeDbMongo             ResourceType = "db.mongo"
+	ResourceTypeDbCosmos            ResourceType = "db.cosmos"
 	ResourceTypeHostContainerApp    ResourceType = "host.containerapp"
 	ResourceTypeOpenAiModel         ResourceType = "ai.openai.model"
 	ResourceTypeMessagingEventHubs  ResourceType = "messaging.eventhubs"
@@ -43,6 +45,8 @@ func (r ResourceType) String() string {
 		return "PostgreSQL"
 	case ResourceTypeDbMongo:
 		return "MongoDB"
+	case ResourceTypeDbCosmos:
+		return "CosmosDB"
 	case ResourceTypeHostContainerApp:
 		return "Container App"
 	case ResourceTypeOpenAiModel:
@@ -96,6 +100,8 @@ func (r *ResourceConfig) MarshalYAML() (interface{}, error) {
 		errMarshal = marshalRawProps(raw.Props.(AIModelProps))
 	case ResourceTypeHostContainerApp:
 		errMarshal = marshalRawProps(raw.Props.(ContainerAppProps))
+	case ResourceTypeDbCosmos:
+		errMarshal = marshalRawProps(raw.Props.(CosmosDBProps))
 	case ResourceTypeMessagingEventHubs:
 		errMarshal = marshalRawProps(raw.Props.(EventHubsProps))
 	case ResourceTypeMessagingServiceBus:
@@ -145,6 +151,12 @@ func (r *ResourceConfig) UnmarshalYAML(value *yaml.Node) error {
 			return err
 		}
 		raw.Props = cap
+	case ResourceTypeDbCosmos:
+		cdp := CosmosDBProps{}
+		if err := unmarshalProps(&cdp); err != nil {
+			return err
+		}
+		raw.Props = cdp
 	case ResourceTypeMessagingEventHubs:
 		ehp := EventHubsProps{}
 		if err := unmarshalProps(&ehp); err != nil {
@@ -191,6 +203,15 @@ type AIModelPropsModel struct {
 	Version string `yaml:"version,omitempty"`
 }
 
+type CosmosDBProps struct {
+	Containers []CosmosDBContainerProps `yaml:"containers,omitempty"`
+}
+
+type CosmosDBContainerProps struct {
+	Name          string   `yaml:"containerName,omitempty"`
+	PartitionKeys []string `yaml:"partitionKeyPaths,omitempty"`
+}
+
 type ServiceBusProps struct {
 	Queues []string `yaml:"queues,omitempty"`
 	Topics []string `yaml:"topics,omitempty"`

cli/azd/pkg/project/scaffold_gen.go

@@ -143,6 +143,19 @@ func infraSpec(projectConfig *ProjectConfig) (*scaffold.InfraSpec, error) {
 			infraSpec.DbCosmosMongo = &scaffold.DatabaseCosmosMongo{
 				DatabaseName: res.Name,
 			}
+		case ResourceTypeDbCosmos:
+			props := res.Props.(CosmosDBProps)
+			containers := make([]scaffold.CosmosSqlDatabaseContainer, 0)
+			for _, c := range props.Containers {
+				containers = append(containers, scaffold.CosmosSqlDatabaseContainer{
+					ContainerName:     c.Name,
+					PartitionKeyPaths: c.PartitionKeys,
+				})
+			}
+			infraSpec.DbCosmos = &scaffold.DatabaseCosmos{
+				DatabaseName: res.Name,
+				Containers:   containers,
+			}
 		case ResourceTypeDbPostgres:
 			infraSpec.DbPostgres = &scaffold.DatabasePostgres{
 				DatabaseName: res.Name,
@@ -285,6 +298,8 @@ func mapHostUses(
 		switch useRes.Type {
 		case ResourceTypeDbMongo:
 			svcSpec.DbCosmosMongo = &scaffold.DatabaseReference{DatabaseName: useRes.Name}
+		case ResourceTypeDbCosmos:
+			svcSpec.DbCosmos = &scaffold.DatabaseReference{DatabaseName: useRes.Name}
 		case ResourceTypeDbPostgres:
 			svcSpec.DbPostgres = &scaffold.DatabaseReference{DatabaseName: useRes.Name}
 		case ResourceTypeDbRedis:

cli/azd/resources/scaffold/templates/main.bicept

@@ -67,6 +67,9 @@ output AZURE_RESOURCE_REDIS_ID string = resources.outputs.AZURE_RESOURCE_REDIS_I
 {{- if .DbPostgres}}
 output AZURE_RESOURCE_{{alphaSnakeUpper .DbPostgres.DatabaseName}}_ID string = resources.outputs.AZURE_RESOURCE_{{alphaSnakeUpper .DbPostgres.DatabaseName}}_ID
 {{- end}}
+{{- if .DbCosmos }}
+output AZURE_RESOURCE_{{alphaSnakeUpper .DbCosmos.DatabaseName}}_ID string = resources.outputs.AZURE_RESOURCE_{{alphaSnakeUpper .DbCosmos.DatabaseName}}_ID
+{{- end}}
 {{- if .StorageAccount }}
 output AZURE_RESOURCE_STORAGE_ID string = resources.outputs.AZURE_RESOURCE_STORAGE_ID
 {{- end}}

cli/azd/resources/scaffold/templates/resources.bicept

@@ -65,8 +65,8 @@ module containerAppsEnvironment 'br/public:avm/res/app/managed-environment:0.4.5
 {{- end}}
 
 {{- if .DbCosmosMongo}}
-module cosmos 'br/public:avm/res/document-db/database-account:0.8.1' = {
-  name: 'cosmos'
+module cosmosMongo 'br/public:avm/res/document-db/database-account:0.8.1' = {
+  name: 'cosmosMongo'
   params: {
     name: '${abbrs.documentDBDatabaseAccounts}${resourceToken}'
     location: location
@@ -101,6 +101,58 @@ module cosmos 'br/public:avm/res/document-db/database-account:0.8.1' = {
 }
 {{- end}}
 
+{{- if .DbCosmos }}
+module cosmos 'br/public:avm/res/document-db/database-account:0.8.1' = {
+  name: 'cosmos'
+  params: {
+    name: '${abbrs.documentDBDatabaseAccounts}${resourceToken}'
+    tags: tags
+    location: location
+    locations: [
+      {
+        failoverPriority: 0
+        isZoneRedundant: false
+        locationName: location
+      }
+    ]
+    networkRestrictions: {
+      ipRules: []
+      virtualNetworkRules: []
+      publicNetworkAccess: 'Enabled'
+    }
+    sqlDatabases: [
+      {
+        name: '{{ .DbCosmos.DatabaseName }}'
+        containers: [
+          {{- range .DbCosmos.Containers}}
+          {
+            name: '{{ .ContainerName }}'
+            paths: [
+              {{- range $path := .PartitionKeyPaths}}
+              '{{ $path }}'
+              {{- end}}
+            ]
+          }
+          {{- end}}
+        ]
+      }
+    ]
+    sqlRoleAssignmentsPrincipalIds: [
+      {{- range .Services}}
+      {{bicepName .Name}}Identity.outputs.principalId
+      {{- end}}
+      principalId
+    ]
+    sqlRoleDefinitions: [
+      {
+        name: 'service-access-cosmos-sql-role'
+      }
+    ]
+    capabilitiesToAdd: [ 'EnableServerless' ]
+  }
+}
+{{- end}}
+
 {{- if .DbPostgres}}
 var databaseName = '{{ .DbPostgres.DatabaseName }}'
 var databaseUser = 'psqladmin'
@@ -353,7 +405,7 @@ module {{bicepName .Name}} 'br/public:avm/res/app/container-app:0.8.0' = {
         {
           name: 'mongodb-url'
           identity:{{bicepName .Name}}Identity.outputs.resourceId
-          keyVaultUrl: cosmos.outputs.exportedSecrets['MONGODB-URL'].secretUri
+          keyVaultUrl: cosmosMongo.outputs.exportedSecrets['MONGODB-URL'].secretUri
         }
         {{- end}}
         {{- if .DbPostgres}}
@@ -407,6 +459,12 @@ module {{bicepName .Name}} 'br/public:avm/res/app/container-app:0.8.0' = {
             secretRef: 'mongodb-url'
           }
           {{- end}}
+          {{- if .DbCosmos}}
+          {
+            name: 'AZURE_COSMOS_ENDPOINT'
+            value: cosmos.outputs.endpoint
+          }
+          {{- end}}
           {{- if .DbPostgres}}
           {
             name: 'POSTGRES_HOST'
@@ -606,6 +664,9 @@ output AZURE_RESOURCE_REDIS_ID string = redis.outputs.resourceId
 {{- if .DbPostgres}}
 output AZURE_RESOURCE_{{alphaSnakeUpper .DbPostgres.DatabaseName}}_ID string = '${postgreServer.outputs.resourceId}/databases/{{.DbPostgres.DatabaseName}}'
 {{- end}}
+{{- if .DbCosmos }}
+output AZURE_RESOURCE_{{alphaSnakeUpper .DbCosmos.DatabaseName}}_ID string = cosmos.outputs.resourceId
+{{- end}}
 {{- if .StorageAccount }}
 output AZURE_RESOURCE_STORAGE_ID string = storageAccount.outputs.resourceId
 {{- end}}

schemas/alpha/azure.yaml.json

@@ -368,6 +368,7 @@
                             "db.postgres",
                             "db.redis",
                             "db.mongo",
+                            "db.cosmos",
                             "ai.openai.model",
                             "host.containerapp",
                             "messaging.eventhubs",
@@ -390,6 +391,7 @@
                     { "if": { "properties": { "type": { "const": "db.postgres"  }}}, "then": { "$ref": "#/definitions/resource"} },
                     { "if": { "properties": { "type": { "const": "db.redis"  }}}, "then": { "$ref": "#/definitions/resource"} },
                     { "if": { "properties": { "type": { "const": "db.mongo"  }}}, "then": { "$ref": "#/definitions/resource"} },
+                    { "if": { "properties": { "type": { "const": "db.cosmos" }}}, "then": { "$ref": "#/definitions/cosmosDbResource"} },
                     { "if": { "properties": { "type": { "const": "messaging.eventhubs" }}}, "then": { "$ref": "#/definitions/eventHubsResource" } },
                     { "if": { "properties": { "type": { "const": "messaging.servicebus" }}}, "then": { "$ref": "#/definitions/serviceBusResource" } },
                     { "if": { "properties": { "type": { "const": "storage"  }}}, "then": { "$ref": "#/definitions/storageAccountResource"} }
@@ -1321,6 +1323,39 @@
                 }
             }
         },
+        "cosmosDbResource": {
+            "type": "object",
+            "description": "A deployed, ready-to-use Azure Cosmos DB for NoSQL.",
+            "additionalProperties": false,
+            "properties": {
+                "type": true,
+                "uses": true,
+                "containers": {
+                    "type": "array",
+                    "title": "Azure Cosmos DB Containers",
+                    "description": "A list of containers in the Azure CosmosDB.",
+                    "items": {
+                        "type": "object",
+                        "additionalProperties": false,
+                        "properties": {
+                            "containerName": {
+                                "type": "string",
+                                "title": "Container Name",
+                                "description": "The name of the container."
+                            },
+                            "partitionKeyPaths": {
+                                "type": "array",
+                                "title": "Partition Key Paths",
+                                "description": "A list of partition key paths for the container.",
+                                "items": {
+                                    "type": "string"
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        },
         "eventHubsResource": {
             "type": "object",
             "description": "An Azure Event Hubs namespace.",