fix: Enhance CosmosDBClient and DatabaseFactory for error handling and concurrency

infra/main.json

@@ -5,8 +5,8 @@
   "metadata": {
     "_generator": {
       "name": "bicep",
-      "version": "0.41.2.15936",
-      "templateHash": "3093757051086668797"
+      "version": "0.42.1.51946",
+      "templateHash": "7222423000870488333"
     },
     "name": "Modernize Your Code Solution Accelerator",
     "description": "CSA CTO Gold Standard Solution Accelerator for Modernize Your Code. \r\n"
@@ -5052,8 +5052,8 @@
           "metadata": {
             "_generator": {
               "name": "bicep",
-              "version": "0.41.2.15936",
-              "templateHash": "8663094775498995429"
+              "version": "0.42.1.51946",
+              "templateHash": "3406526791248457038"
             }
           },
           "definitions": {
@@ -12895,11 +12895,11 @@
       },
       "dependsOn": [
         "applicationInsights",
-        "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').storageBlob)]",
         "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').ods)]",
-        "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').agentSvc)]",
         "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').monitor)]",
         "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').oms)]",
+        "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').storageBlob)]",
+        "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').agentSvc)]",
         "dataCollectionEndpoint",
         "logAnalyticsWorkspace",
         "virtualNetwork"
@@ -25611,8 +25611,8 @@
           "metadata": {
             "_generator": {
               "name": "bicep",
-              "version": "0.41.2.15936",
-              "templateHash": "17285204072656433491"
+              "version": "0.42.1.51946",
+              "templateHash": "16969185198334420434"
             },
             "name": "AI Services and Project Module",
             "description": "This module creates an AI Services resource and an AI Foundry project within it. It supports private networking, OpenAI deployments, and role assignments."
@@ -26952,8 +26952,8 @@
                   "metadata": {
                     "_generator": {
                       "name": "bicep",
-                      "version": "0.41.2.15936",
-                      "templateHash": "5121330425393020264"
+                      "version": "0.42.1.51946",
+                      "templateHash": "4140498216793917924"
                     }
                   },
                   "definitions": {
@@ -27910,7 +27910,10 @@
                         "raiPolicyName": "[tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'raiPolicyName')]",
                         "versionUpgradeOption": "[tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'versionUpgradeOption')]"
                       },
-                      "sku": "[coalesce(tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'sku'), createObject('name', parameters('sku'), 'capacity', tryGet(parameters('sku'), 'capacity'), 'tier', tryGet(parameters('sku'), 'tier'), 'size', tryGet(parameters('sku'), 'size'), 'family', tryGet(parameters('sku'), 'family')))]"
+                      "sku": "[coalesce(tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'sku'), createObject('name', parameters('sku'), 'capacity', tryGet(parameters('sku'), 'capacity'), 'tier', tryGet(parameters('sku'), 'tier'), 'size', tryGet(parameters('sku'), 'size'), 'family', tryGet(parameters('sku'), 'family')))]",
+                      "dependsOn": [
+                        "aiProject"
+                      ]
                     },
                     "cognitiveService_lock": {
                       "condition": "[and(not(empty(coalesce(parameters('lock'), createObject()))), not(equals(tryGet(parameters('lock'), 'kind'), 'None')))]",
@@ -28664,8 +28667,8 @@
                           "metadata": {
                             "_generator": {
                               "name": "bicep",
-                              "version": "0.41.2.15936",
-                              "templateHash": "10989408486030617267"
+                              "version": "0.42.1.51946",
+                              "templateHash": "2422737205646151487"
                             }
                           },
                           "definitions": {
@@ -28818,8 +28821,8 @@
                           "metadata": {
                             "_generator": {
                               "name": "bicep",
-                              "version": "0.41.2.15936",
-                              "templateHash": "7933643033523871028"
+                              "version": "0.42.1.51946",
+                              "templateHash": "11911242767938607365"
                             }
                           },
                           "definitions": {
@@ -29036,8 +29039,8 @@
                   "metadata": {
                     "_generator": {
                       "name": "bicep",
-                      "version": "0.41.2.15936",
-                      "templateHash": "5121330425393020264"
+                      "version": "0.42.1.51946",
+                      "templateHash": "4140498216793917924"
                     }
                   },
                   "definitions": {
@@ -29994,7 +29997,10 @@
                         "raiPolicyName": "[tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'raiPolicyName')]",
                         "versionUpgradeOption": "[tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'versionUpgradeOption')]"
                       },
-                      "sku": "[coalesce(tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'sku'), createObject('name', parameters('sku'), 'capacity', tryGet(parameters('sku'), 'capacity'), 'tier', tryGet(parameters('sku'), 'tier'), 'size', tryGet(parameters('sku'), 'size'), 'family', tryGet(parameters('sku'), 'family')))]"
+                      "sku": "[coalesce(tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'sku'), createObject('name', parameters('sku'), 'capacity', tryGet(parameters('sku'), 'capacity'), 'tier', tryGet(parameters('sku'), 'tier'), 'size', tryGet(parameters('sku'), 'size'), 'family', tryGet(parameters('sku'), 'family')))]",
+                      "dependsOn": [
+                        "aiProject"
+                      ]
                     },
                     "cognitiveService_lock": {
                       "condition": "[and(not(empty(coalesce(parameters('lock'), createObject()))), not(equals(tryGet(parameters('lock'), 'kind'), 'None')))]",
@@ -30748,8 +30754,8 @@
                           "metadata": {
                             "_generator": {
                               "name": "bicep",
-                              "version": "0.41.2.15936",
-                              "templateHash": "10989408486030617267"
+                              "version": "0.42.1.51946",
+                              "templateHash": "2422737205646151487"
                             }
                           },
                           "definitions": {
@@ -30902,8 +30908,8 @@
                           "metadata": {
                             "_generator": {
                               "name": "bicep",
-                              "version": "0.41.2.15936",
-                              "templateHash": "7933643033523871028"
+                              "version": "0.42.1.51946",
+                              "templateHash": "11911242767938607365"
                             }
                           },
                           "definitions": {
@@ -31917,8 +31923,8 @@
       "dependsOn": [
         "aiServices",
         "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').openAI)]",
-        "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').cognitiveServices)]",
         "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').aiServices)]",
+        "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').cognitiveServices)]",
         "virtualNetwork"
       ]
     },
@@ -31974,8 +31980,8 @@
           "metadata": {
             "_generator": {
               "name": "bicep",
-              "version": "0.41.2.15936",
-              "templateHash": "3881415837167031634"
+              "version": "0.42.1.51946",
+              "templateHash": "522477461329004641"
             }
           },
           "definitions": {
@@ -40219,8 +40225,8 @@
           "metadata": {
             "_generator": {
               "name": "bicep",
-              "version": "0.41.2.15936",
-              "templateHash": "15962472891869337617"
+              "version": "0.42.1.51946",
+              "templateHash": "15355322017409205910"
             }
           },
           "definitions": {
@@ -44082,8 +44088,8 @@
           "metadata": {
             "_generator": {
               "name": "bicep",
-              "version": "0.41.2.15936",
-              "templateHash": "17636459140972536078"
+              "version": "0.42.1.51946",
+              "templateHash": "4242598725709304634"
             }
           },
           "definitions": {

infra/modules/ai-foundry/dependencies.bicep

@@ -208,6 +208,9 @@ resource cognitiveService_deployments 'Microsoft.CognitiveServices/accounts/depl
       size: sku.?size
       family: sku.?family
     }
+    dependsOn: [
+      aiProject
+    ]
   }
 ]
 

src/backend/common/database/cosmosdb.py

@@ -1,11 +1,13 @@
+import asyncio
 from datetime import datetime, timezone
 from typing import Dict, List, Optional
 from uuid import UUID, uuid4
 
 from azure.cosmos.aio import CosmosClient
 from azure.cosmos.aio._database import DatabaseProxy
 from azure.cosmos.exceptions import (
-    CosmosResourceExistsError
+    CosmosResourceExistsError,
+    CosmosResourceNotFoundError,
 )
 
 from common.database.database_base import DatabaseBase
@@ -85,9 +87,26 @@ async def create_batch(self, user_id: str, batch_id: UUID) -> BatchRecord:
                 await self.batch_container.create_item(body=batch.dict())
                 return batch
             except CosmosResourceExistsError:
-                self.logger.info(f"Batch with ID {batch_id} already exists")
-                batchexists = await self.get_batch(user_id, str(batch_id))
-                return batchexists
+                self.logger.info("Batch already exists, reading existing record", batch_id=str(batch_id))
+                # Retry read with backoff to handle replication lag after 409 conflict
+                for attempt in range(3):
+                    try:
+                        batchexists = await self.batch_container.read_item(
+                            item=str(batch_id), partition_key=str(batch_id)
+                        )
+                        if batchexists.get("user_id") != user_id:
+                            self.logger.error("Batch belongs to a different user", batch_id=str(batch_id))
+                            raise CosmosResourceNotFoundError(message="Batch not found")
+                        self.logger.info("Returning existing batch record", batch_id=str(batch_id))
+                        return BatchRecord.fromdb(batchexists)
+                    except CosmosResourceNotFoundError:
+                        if attempt < 2:
+                            self.logger.info("Batch read returned 404 after conflict, retrying", batch_id=str(batch_id), attempt=attempt + 1)
+                            await asyncio.sleep(0.5 * (attempt + 1))
+                        else:
+                            raise RuntimeError(
+                                f"Batch {batch_id} already exists but could not be read after retries"
+                            )
 
         except Exception as e:
             self.logger.error("Failed to create batch", error=str(e))
@@ -158,7 +177,7 @@ async def get_batch(self, user_id: str, batch_id: str) -> Optional[Dict]:
             ]
             batch = None
             async for item in self.batch_container.query_items(
-                query=query, parameters=params
+                query=query, parameters=params, partition_key=batch_id
             ):
                 batch = item
 
@@ -173,7 +192,7 @@ async def get_file(self, file_id: str) -> Optional[Dict]:
             params = [{"name": "@file_id", "value": file_id}]
             file_entry = None
             async for item in self.file_container.query_items(
-                query=query, parameters=params
+                query=query, parameters=params, partition_key=file_id
             ):
                 file_entry = item
             return file_entry
@@ -209,7 +228,7 @@ async def get_batch_from_id(self, batch_id: str) -> Dict:
 
             batch = None  # Store the batch
             async for item in self.batch_container.query_items(
-                query=query, parameters=params
+                query=query, parameters=params, partition_key=batch_id
             ):
                 batch = item  # Assign the batch to the variable
 
@@ -335,11 +354,14 @@ async def add_file_log(
             raise
 
     async def update_batch_entry(
-        self, batch_id: str, user_id: str, status: ProcessStatus, file_count: int
+        self, batch_id: str, user_id: str, status: ProcessStatus, file_count: int,
+        existing_batch: Optional[Dict] = None
     ):
-        """Update batch status."""
+        """Update batch status. If existing_batch is provided, skip the re-fetch."""
         try:
-            batch = await self.get_batch(user_id, batch_id)
+            batch = existing_batch
+            if batch is None:
+                batch = await self.get_batch(user_id, batch_id)
             if not batch:
                 raise ValueError("Batch not found")
 

src/backend/common/database/database_factory.py

@@ -9,25 +9,40 @@
 
 class DatabaseFactory:
     _instance: Optional[DatabaseBase] = None
+    _lock: Optional[asyncio.Lock] = None
     _logger = AppLogger("DatabaseFactory")
 
+    @staticmethod
+    def _get_lock() -> asyncio.Lock:
+        if DatabaseFactory._lock is None:
+            DatabaseFactory._lock = asyncio.Lock()
+        return DatabaseFactory._lock
+
     @staticmethod
     async def get_database():
+        if DatabaseFactory._instance is not None:
+            return DatabaseFactory._instance
+
+        async with DatabaseFactory._get_lock():
+            # Double-check after acquiring the lock
+            if DatabaseFactory._instance is not None:
+                return DatabaseFactory._instance
 
-        config = Config()  # Create an instance of Config
+            config = Config()  # Create an instance of Config
 
-        cosmos_db_client = CosmosDBClient(
-            endpoint=config.cosmosdb_endpoint,
-            credential=config.get_azure_credentials(),
-            database_name=config.cosmosdb_database,
-            batch_container=config.cosmosdb_batch_container,
-            file_container=config.cosmosdb_file_container,
-            log_container=config.cosmosdb_log_container,
-        )
+            cosmos_db_client = CosmosDBClient(
+                endpoint=config.cosmosdb_endpoint,
+                credential=config.get_azure_credentials(),
+                database_name=config.cosmosdb_database,
+                batch_container=config.cosmosdb_batch_container,
+                file_container=config.cosmosdb_file_container,
+                log_container=config.cosmosdb_log_container,
+            )
 
-        await cosmos_db_client.initialize_cosmos()
+            await cosmos_db_client.initialize_cosmos()
 
-        return cosmos_db_client
+            DatabaseFactory._instance = cosmos_db_client
+            return cosmos_db_client
 
 
 # Local testing of config and code

src/backend/common/services/batch_service.py

@@ -288,8 +288,9 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi
             self.logger.info("File uploaded to blob storage", filename=file.filename, batch_id=batch_id)
 
             # Create file entry
-            await self.database.add_file(batch_id, file_id, file.filename, blob_path)
-            file_record = await self.database.get_file(file_id)
+            file_record_obj = await self.database.add_file(batch_id, file_id, file.filename, blob_path)
+            file_record_dict = getattr(file_record_obj, "dict", None)
+            file_record = file_record_dict() if callable(file_record_dict) else file_record_obj
 
             await self.database.add_file_log(
                 UUID(file_id),
@@ -308,6 +309,7 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi
                     user_id,
                     ProcessStatus.READY_TO_PROCESS,
                     batch["file_count"],
+                    existing_batch=batch,
                 )
                 # Return response
                 return {"batch": batch, "file": file_record}
@@ -318,7 +320,8 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi
                 batch.file_count = len(files)
                 batch.updated_at = datetime.utcnow().isoformat()
                 await self.database.update_batch_entry(
-                    batch_id, user_id, ProcessStatus.READY_TO_PROCESS, batch.file_count
+                    batch_id, user_id, ProcessStatus.READY_TO_PROCESS, batch.file_count,
+                    existing_batch=batch.dict(),
                 )
                 # Return response
                 return {"batch": batch, "file": file_record}