Merge pull request #56 from ks6088ts-labs/feature/issue-48_support-cosmosdb

support cosmosdb

Author: ks6088ts

azure_cosmos_db.env.sample

@@ -0,0 +1,3 @@
+AZURE_COSMOS_DB_ENDPOINT = "<cosmosdb-endpoint>"
+AZURE_COSMOS_DB_KEY = "<cosmosdb-key>"
+AZURE_COSMOS_DB_CONNECTION_STRING = "<cosmosdb-connection-string>"

backend/fastapi.py

@@ -3,6 +3,7 @@
 
 from backend.routers import azure_ai_document_intelligence as azure_ai_document_intelligence_router
 from backend.routers import azure_ai_vision as azure_ai_vision_router
+from backend.routers import azure_cosmos_db as azure_cosmos_db_router
 from backend.routers import azure_event_grid as azure_event_grid_router
 from backend.routers import azure_openai as azure_openai_router
 from backend.routers import azure_storage_blob as azure_storage_blob_router
@@ -19,6 +20,7 @@
     azure_openai_router.router,
     azure_storage_blob_router.router,
     azure_storage_queue_router.router,
+    azure_cosmos_db_router.router,
 ]:
     app.include_router(router)
 

backend/internals/azure_cosmos_db.py

@@ -0,0 +1,91 @@
+from logging import getLogger
+
+from azure.core.exceptions import AzureError
+from azure.cosmos import CosmosClient, PartitionKey
+from azure.cosmos.container import ContainerProxy
+from azure.cosmos.database import DatabaseProxy
+
+from backend.settings.azure_cosmos_db import Settings
+
+logger = getLogger(__name__)
+
+
+class Client:
+    def __init__(self, settings: Settings) -> None:
+        self.settings = settings
+
+    def get_client(self) -> CosmosClient:
+        return CosmosClient.from_connection_string(
+            conn_str=self.settings.azure_cosmos_db_connection_string,
+        )
+
+    def get_database(self, database_id: str) -> DatabaseProxy:
+        return self.get_client().get_database_client(database_id)
+
+    def get_container(
+        self,
+        database_id: str,
+        container_id: str,
+    ) -> ContainerProxy:
+        return self.get_database(database_id=database_id).get_container_client(container=container_id)
+
+    def create_database(
+        self,
+        database_id: str,
+    ) -> str:
+        try:
+            self.get_client().create_database_if_not_exists(
+                id=database_id,
+            )
+        except AzureError as e:
+            logger.error(e)
+        return database_id
+
+    def create_container(
+        self,
+        database_id: str,
+        container_id: str,
+        partition_key_path="/id",
+    ) -> str:
+        try:
+            self.get_database(database_id=database_id).create_container_if_not_exists(
+                id=container_id,
+                partition_key=PartitionKey(
+                    path=partition_key_path,
+                ),
+            )
+        except AzureError as e:
+            logger.error(e)
+        return container_id
+
+    def create_item(
+        self,
+        container: ContainerProxy,
+        item: dict,
+    ) -> dict:
+        return container.create_item(
+            body=item,
+        )
+
+    def read_item(
+        self,
+        container: ContainerProxy,
+        item_id: str,
+    ) -> dict:
+        return container.read_item(
+            item=item_id,
+            partition_key=item_id,
+        )
+
+    def query_items(
+        self,
+        container: ContainerProxy,
+        query: str,
+        parameters: list | None = None,
+        enable_cross_partition_query: bool = True,
+    ) -> list:
+        return container.query_items(
+            query=query,
+            parameters=parameters,
+            enable_cross_partition_query=enable_cross_partition_query,
+        )

backend/routers/azure_cosmos_db.py

@@ -0,0 +1,69 @@
+from logging import getLogger
+
+from fastapi import APIRouter
+
+from backend.internals.azure_cosmos_db import Client
+from backend.schemas import azure_cosmos_db as azure_cosmos_db_schemas
+from backend.settings.azure_cosmos_db import Settings
+
+logger = getLogger(__name__)
+
+client = Client(
+    settings=Settings(),
+)
+
+router = APIRouter(
+    prefix="/azure_cosmos_db",
+    tags=["azure_cosmos_db"],
+    responses={404: {"description": "Not found"}},
+)
+
+
+@router.post(
+    "/database",
+    response_model=azure_cosmos_db_schemas.CreateDatabaseResponse,
+    status_code=200,
+)
+async def create_database(body: azure_cosmos_db_schemas.CreateDatabaseRequest):
+    database_id = client.create_database(
+        database_id=body.database_id,
+    )
+    return azure_cosmos_db_schemas.CreateDatabaseResponse(
+        database_id=database_id,
+    )
+
+
+@router.post(
+    "/container",
+    response_model=azure_cosmos_db_schemas.CreateContainerResponse,
+    status_code=200,
+)
+async def create_container(body: azure_cosmos_db_schemas.CreateContainerRequest):
+    container_id = client.create_container(
+        container_id=body.container_id,
+        database_id=body.database_id,
+    )
+    return azure_cosmos_db_schemas.CreateContainerResponse(
+        container_id=container_id,
+    )
+
+
+@router.post(
+    "/item",
+    response_model=azure_cosmos_db_schemas.CreateItemResponse,
+    status_code=200,
+)
+async def create_item(body: azure_cosmos_db_schemas.CreateItemRequest):
+    container = client.get_container(
+        container_id=body.container_id,
+        database_id=body.database_id,
+    )
+    created_item = client.create_item(
+        container=container,
+        item=body.item,
+    )
+    return azure_cosmos_db_schemas.CreateItemResponse(
+        container_id=body.container_id,
+        database_id=body.database_id,
+        item=created_item,
+    )

backend/schemas/azure_cosmos_db.py

@@ -0,0 +1,34 @@
+from logging import getLogger
+
+from pydantic import BaseModel
+
+logger = getLogger(__name__)
+
+
+class CreateDatabaseRequest(BaseModel):
+    database_id: str
+
+
+class CreateDatabaseResponse(BaseModel):
+    database_id: str
+
+
+class CreateContainerRequest(BaseModel):
+    container_id: str
+    database_id: str
+
+
+class CreateContainerResponse(BaseModel):
+    container_id: str
+
+
+class CreateItemRequest(BaseModel):
+    container_id: str
+    database_id: str
+    item: dict
+
+
+class CreateItemResponse(BaseModel):
+    container_id: str
+    database_id: str
+    item: dict

backend/settings/azure_cosmos_db.py

@@ -0,0 +1,13 @@
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Settings(BaseSettings):
+    azure_cosmos_db_endpoint: str = "https://<your-cosmos-db-name>.documents.azure.com:443/"
+    azure_cosmos_db_key: str = "<your-cosmos-db-key>"
+    azure_cosmos_db_connection_string: str = "<your-cosmos-db-connection-string>"
+
+    model_config = SettingsConfigDict(
+        env_file="azure_cosmos_db.env",
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )

docs/README.md

@@ -60,3 +60,11 @@
 ### Azure AI Document Intelligence
 
 - [Azure AI Document Intelligence client library for Python](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/documentintelligence/azure-ai-documentintelligence/README.md)
+
+### Azure Cosmos DB
+
+- [Manage Azure Cosmos DB for NoSQL resources with Bicep > Azure Cosmos DB account with autoscale throughput](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/manage-with-bicep#azure-cosmos-db-account-with-autoscale-throughput)
+- [Get started with Azure Cosmos DB for NoSQL using Python](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-python-get-started?tabs=env-virtual%2Cazure-cli%2Clinux)
+- [Create a database in Azure Cosmos DB for NoSQL using Python](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-python-create-database)
+- [Create a container in Azure Cosmos DB for NoSQL using Python](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-python-create-container)
+- [Sample - demonstrates the basic CRUD operations on a Item resource for Azure Cosmos](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/samples/document_management.py)

infra/main.bicep

@@ -8,7 +8,6 @@ param location string = resourceGroup().location
 @description('Specifies the resource tags.')
 param tags object = {}
 
-// OpenAI
 @description('Specifies the name of the Azure OpenAI resource.')
 param openAiName string = '${prefix}openai'
 
@@ -24,6 +23,9 @@ param storageAccountName string = '${prefix}sa'
 @description('Specifies the name of the Azure Event Grid resource.')
 param eventGridName string = '${prefix}eg'
 
+@description('Specifies the name of the Azure Cosmos DB resource.')
+param cosmosDbName string = '${prefix}cosmosdb'
+
 module openAi './modules/openAi.bicep' = {
   name: 'openAi'
   params: {
@@ -73,6 +75,17 @@ module eventGrid './modules/eventGrid.bicep' = {
   }
 }
 
+module cosmosDb './modules/cosmosDb.bicep' = {
+  name: 'cosmosDb'
+  params: {
+    name: cosmosDbName
+    location: location
+    tags: tags
+    primaryRegion: location
+    secondaryRegion: 'japanwest'
+  }
+}
+
 // Output
 output cognitiveServicesName string = cognitiveServices.outputs.name
 output cognitiveServicesEndpoint string = cognitiveServices.outputs.endpoint

infra/modules/cosmosDb.bicep

@@ -0,0 +1,91 @@
+// Parameters
+@description('Specifies the name of the Azure Cosmos DB resource')
+param name string
+
+@description('Specifies the primary location of Azure resources.')
+param location string = resourceGroup().location
+
+@description('Specifies the resource tags.')
+param tags object = {}
+
+@description('Specifies the name of the Azure Cosmos DB account name, max length 44 characters, lowercase.')
+param accountName string = toLower('${name}account')
+
+@description('The primary region for the Cosmos DB account.')
+param primaryRegion string
+
+@description('The secondary region for the Cosmos DB account.')
+param secondaryRegion string
+
+@description('The default consistency level of the Cosmos DB account.')
+@allowed([
+  'Eventual'
+  'ConsistentPrefix'
+  'Session'
+  'BoundedStaleness'
+  'Strong'
+])
+param defaultConsistencyLevel string = 'Session'
+
+@description('Max stale requests. Required for BoundedStaleness. Valid ranges, Single Region: 10 to 2147483647. Multi Region: 100000 to 2147483647.')
+@minValue(10)
+@maxValue(2147483647)
+param maxStalenessPrefix int = 100000
+
+@description('Max lag time (minutes). Required for BoundedStaleness. Valid ranges, Single Region: 5 to 84600. Multi Region: 300 to 86400.')
+@minValue(5)
+@maxValue(86400)
+param maxIntervalInSeconds int = 300
+
+@description('Enable system managed failover for regions')
+param systemManagedFailover bool = true
+
+var consistencyPolicy = {
+  Eventual: {
+    defaultConsistencyLevel: 'Eventual'
+  }
+  ConsistentPrefix: {
+    defaultConsistencyLevel: 'ConsistentPrefix'
+  }
+  Session: {
+    defaultConsistencyLevel: 'Session'
+  }
+  BoundedStaleness: {
+    defaultConsistencyLevel: 'BoundedStaleness'
+    maxStalenessPrefix: maxStalenessPrefix
+    maxIntervalInSeconds: maxIntervalInSeconds
+  }
+  Strong: {
+    defaultConsistencyLevel: 'Strong'
+  }
+}
+var locations = [
+  {
+    locationName: primaryRegion
+    failoverPriority: 0
+    isZoneRedundant: false
+  }
+  {
+    locationName: secondaryRegion
+    failoverPriority: 1
+    isZoneRedundant: false
+  }
+]
+
+resource account 'Microsoft.DocumentDB/databaseAccounts@2024-02-15-preview' = {
+  name: accountName
+  location: location
+  tags: tags
+  kind: 'GlobalDocumentDB'
+  properties: {
+    consistencyPolicy: consistencyPolicy[defaultConsistencyLevel]
+    locations: locations
+    databaseAccountOfferType: 'Standard'
+    enableAutomaticFailover: systemManagedFailover
+  }
+}
+
+// Outputs
+output accountId string = account.id
+output accountName string = account.name
+output accountEndpoint string = account.properties.documentEndpoint