Enhance CosmosDBClient and DatabaseFactory for improved error handling and concurrency

- Added asyncio support and a lock mechanism in DatabaseFactory to ensure thread safety.
- Implemented retry logic with backoff for reading existing batches in CosmosDBClient.
- Updated batch service to handle existing batch entries more efficiently.

Author: Pavan-Microsoft

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,21 @@ 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:
+                            raise ValueError("Batch belongs to a different user")
+                        return batchexists
+                    except CosmosResourceNotFoundError:
+                        if attempt < 2:
+                            await asyncio.sleep(0.5 * (attempt + 1))
+                        else:
+                            return batch
 
         except Exception as e:
             self.logger.error("Failed to create batch", error=str(e))
@@ -158,7 +172,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 +187,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 +223,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 +349,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,34 @@
 
 class DatabaseFactory:
     _instance: Optional[DatabaseBase] = None
+    _lock: asyncio.Lock = asyncio.Lock()
     _logger = AppLogger("DatabaseFactory")
 
     @staticmethod
     async def get_database():
+        if DatabaseFactory._instance is not None:
+            return DatabaseFactory._instance
 
-        config = Config()  # Create an instance of Config
+        async with DatabaseFactory._lock:
+            # Double-check after acquiring the lock
+            if DatabaseFactory._instance is not None:
+                return DatabaseFactory._instance
 
-        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,
-        )
+            config = Config()  # Create an instance of Config
 
-        await cosmos_db_client.initialize_cosmos()
+            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,
+            )
 
-        return cosmos_db_client
+            await cosmos_db_client.initialize_cosmos()
+
+            DatabaseFactory._instance = cosmos_db_client
+            return cosmos_db_client
 
 
 # Local testing of config and code

src/backend/common/services/batch_service.py

@@ -287,8 +287,8 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi
             )
 
             # 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 = file_record_obj.dict() if hasattr(file_record_obj, 'dict') else file_record_obj
 
             await self.database.add_file_log(
                 UUID(file_id),
@@ -307,6 +307,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}
@@ -317,7 +318,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}