implement cosmos db APIs

Author: ks6088ts

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

@@ -22,23 +22,31 @@ def get_client(self) -> CosmosClient:
     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:
+    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) -> None:
+    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",
-    ) -> None:
+    ) -> str:
         try:
             self.get_database(database_id=database_id).create_container_if_not_exists(
                 id=container_id,
@@ -48,6 +56,7 @@ def create_container(
             )
         except AzureError as e:
             logger.error(e)
+        return container_id
 
     def create_item(
         self,

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

tests/test_smoke_azure_cosmos_db.py

@@ -0,0 +1,48 @@
+from logging import getLogger
+from uuid import uuid4
+
+import pytest
+
+from tests.utilities import RUN_TEST, client
+
+logger = getLogger(__name__)
+
+
+@pytest.mark.skipif(RUN_TEST, reason="need to launch the backend server first")
+def test_azure_cosmos_db():
+    path_format = "/azure_cosmos_db/{0}"
+    database_id = str(uuid4())
+    container_id = str(uuid4())
+
+    response = client.post(
+        url=path_format.format("database"),
+        json={
+            "database_id": database_id,
+        },
+    )
+    assert response.status_code == 200
+    logger.info(f"response: {response.json()}")
+
+    response = client.post(
+        url=path_format.format("container"),
+        json={
+            "database_id": database_id,
+            "container_id": container_id,
+        },
+    )
+    assert response.status_code == 200
+    logger.info(f"response: {response.json()}")
+
+    response = client.post(
+        url=path_format.format("item"),
+        json={
+            "database_id": database_id,
+            "container_id": container_id,
+            "item": {
+                "id": str(uuid4()),
+                "name": "test",
+            },
+        },
+    )
+    assert response.status_code == 200
+    logger.info(f"response: {response.json()}")