"""PostgreSQL storage for extracted records.""" import asyncio import time from dataclasses import dataclass from typing import Any from uuid import UUID, uuid4 import structlog from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.exc import IntegrityError, OperationalError from sqlalchemy.ext.asyncio import AsyncSession from src.core.exceptions import IndexingError from src.db.models import ExtractedRecord as ExtractedRecordModel from src.extraction.models import ExtractedRecord @dataclass class StoreResult: """Result from storing records in PostgreSQL. Attributes: count: Number of records successfully stored record_ids: List of UUIDs for each stored record (in same order as input) """ count: int record_ids: list[UUID] # Constants MAX_BATCH_SIZE = 500 MAX_RETRIES = 3 RETRY_BACKOFF_SECONDS = [1, 2, 4] # Exponential backoff: 1s, 2s, 4s logger = structlog.get_logger() class StructuredStore: """Manages PostgreSQL storage for extracted records. Provides optimized batch insertion of extracted records with: - JSONB storage for content, resolved_content, and headers - Batch processing (500 records per transaction) - ON CONFLICT handling for idempotent operations - Retry logic for transient database errors - Performance logging """ def __init__(self, db_session: AsyncSession, tenant_id: UUID | None = None): """Initialize StructuredStore. Args: db_session: Async SQLAlchemy session for database operations tenant_id: Optional tenant UUID for RLS policy """ self.db = db_session self.tenant_id = tenant_id self.logger = logger.bind(component="structured_store") async def _set_rls_context(self) -> None: """Set RLS tenant context on the session if tenant_id is available.""" if self.tenant_id: from sqlalchemy import text from uuid import UUID as _UUID _tid = str(_UUID(str(self.tenant_id))) # validate UUID format await self.db.execute(text(f"SET LOCAL app.tenant_id = '{_tid}'")) async def store_records( self, records: list[ExtractedRecord], document_id: UUID, ) -> StoreResult: """Store extracted records in PostgreSQL. Processes records in batches of 500, using ON CONFLICT for idempotency. Includes retry logic for transient database errors. Args: records: List of ExtractedRecord from extraction pipeline document_id: UUID of source document Returns: StoreResult with count and list of record UUIDs (for Milvus linking) Raises: IndexingError: If storage fails after retries """ if not records: self.logger.warning("store_records_called_with_empty_list", document_id=str(document_id)) return StoreResult(count=0, record_ids=[]) start_time = time.time() total_stored = 0 all_record_ids: list[UUID] = [] self.logger.info( "store_records_started", document_id=str(document_id), record_count=len(records), ) try: # Split records into batches batches = [ records[i : i + MAX_BATCH_SIZE] for i in range(0, len(records), MAX_BATCH_SIZE) ] for batch_idx, batch in enumerate(batches): batch_start = time.time() # Store batch with retry logic - returns (count, ids) stored_count, batch_ids = await self._store_batch_with_retry( batch, document_id, batch_idx ) total_stored += stored_count all_record_ids.extend(batch_ids) batch_duration = (time.time() - batch_start) * 1000 self.logger.debug( "batch_stored", batch_idx=batch_idx, batch_size=len(batch), stored_count=stored_count, duration_ms=int(batch_duration), ) # Log final metrics total_duration = (time.time() - start_time) * 1000 records_per_second = ( len(records) / (total_duration / 1000) if total_duration > 0 else 0 ) self.logger.info( "store_records_completed", document_id=str(document_id), records_stored=total_stored, record_ids_count=len(all_record_ids), duration_ms=int(total_duration), records_per_second=int(records_per_second), ) return StoreResult(count=total_stored, record_ids=all_record_ids) except Exception as e: duration = (time.time() - start_time) * 1000 self.logger.error( "store_records_failed", document_id=str(document_id), error=str(e), error_type=type(e).__name__, duration_ms=int(duration), ) raise IndexingError( message=f"Failed to store records for document {document_id}", details={ "document_id": str(document_id), "error": str(e), "error_type": type(e).__name__, }, ) from e async def _store_batch_with_retry( self, batch: list[ExtractedRecord], document_id: UUID, batch_idx: int, ) -> tuple[int, list[UUID]]: """Store a batch of records with retry logic for transient errors. Args: batch: Batch of records to store document_id: UUID of source document batch_idx: Index of this batch (for logging) Returns: Tuple of (count of stored records, list of record UUIDs) Raises: IndexingError: If storage fails after all retries """ for attempt in range(MAX_RETRIES): try: return await self._store_batch(batch, document_id) except OperationalError as e: # Transient database error (connection loss, timeout) if attempt < MAX_RETRIES - 1: backoff = RETRY_BACKOFF_SECONDS[attempt] self.logger.warning( "batch_store_retry", batch_idx=batch_idx, attempt=attempt + 1, max_retries=MAX_RETRIES, error=str(e), retry_after_seconds=backoff, ) await asyncio.sleep(backoff) else: # Final attempt failed self.logger.error( "batch_store_failed_after_retries", batch_idx=batch_idx, attempts=MAX_RETRIES, error=str(e), ) raise except IntegrityError as e: # Permanent error (data integrity issue) self.logger.error( "batch_store_integrity_error", batch_idx=batch_idx, error=str(e), ) raise except Exception as e: # Unexpected error self.logger.error( "batch_store_unexpected_error", batch_idx=batch_idx, error=str(e), error_type=type(e).__name__, ) raise # Should never reach here (covered by raises above) return 0, [] async def _store_batch( self, batch: list[ExtractedRecord], document_id: UUID, ) -> tuple[int, list[UUID]]: """Store a single batch of records in a transaction. Args: batch: Batch of records to store document_id: UUID of source document Returns: Tuple of (count of stored records, list of record UUIDs) Raises: SQLAlchemy exceptions on database errors """ # Generate UUIDs upfront so we can return them for Milvus linking record_ids = [uuid4() for _ in batch] # Convert ExtractedRecord objects to database model dicts with pre-generated IDs db_records = [ self._to_db_model(record, document_id, record_id) for record, record_id in zip(batch, record_ids) ] # Build INSERT statement stmt = pg_insert(ExtractedRecordModel).values(db_records) # Execute within transaction (only start new transaction if not already in one) if not self.db.in_transaction(): async with self.db.begin(): await self._set_rls_context() result = await self.db.execute(stmt) return result.rowcount, record_ids else: # Already in a transaction, just execute await self._set_rls_context() result = await self.db.execute(stmt) return result.rowcount, record_ids def _to_db_model( self, extracted_record: ExtractedRecord, document_id: UUID, record_id: UUID, ) -> dict[str, Any]: """Convert ExtractedRecord to database model dict. Args: extracted_record: Extraction pipeline record document_id: UUID of source document record_id: Pre-generated UUID for this record (for Milvus linking) Returns: Dictionary with database model fields """ # Extract source metadata source = extracted_record._source # Build database record with explicit ID for Milvus linking db_record = { "id": record_id, "document_id": document_id, "sheet_name": source["sheet"], "row_number": source["row"], "col_range": source.get("col_range"), "content": extracted_record.content, "headers": extracted_record.headers, } # Add resolved_content if present (from symbol resolution) # Note: ExtractedRecord doesn't have resolved_content attribute yet # This will be added in Story 3.6 (Symbol Resolution) # For now, we check if it exists via getattr resolved_content = getattr(extracted_record, "resolved_content", None) if resolved_content: db_record["resolved_content"] = resolved_content # Add tags if present (from SemanticTagger) # Tags are TagAssignment objects or dicts, serialize to JSON-compatible format tags = getattr(extracted_record, "tags", None) if tags: serialized_tags = [] for tag in tags: if hasattr(tag, "model_dump"): # Pydantic v2 model serialized_tags.append(tag.model_dump()) elif hasattr(tag, "dict"): # Pydantic v1 model serialized_tags.append(tag.dict()) elif isinstance(tag, dict): # Already a dict serialized_tags.append(tag) else: # String or other - skip self.logger.warning( "skipping_invalid_tag_type", tag_type=type(tag).__name__, ) db_record["tags"] = serialized_tags else: db_record["tags"] = [] return db_record async def delete_records_by_document_id(self, document_id: UUID) -> int: """Delete all records for a document using CASCADE. Args: document_id: Document UUID to filter records Returns: Count of deleted records Raises: IndexingError: If deletion fails """ try: from sqlalchemy import delete # Build DELETE statement stmt = delete(ExtractedRecordModel).where( ExtractedRecordModel.document_id == document_id ) # Execute deletion (CASCADE will handle dependencies) if not self.db.in_transaction(): async with self.db.begin(): await self._set_rls_context() result = await self.db.execute(stmt) deleted_count = result.rowcount else: # Already in a transaction await self._set_rls_context() result = await self.db.execute(stmt) deleted_count = result.rowcount self.logger.info( "records_deleted", document_id=str(document_id), deleted_count=deleted_count, ) return deleted_count except Exception as e: self.logger.error( "records_deletion_failed", document_id=str(document_id), error=str(e), error_type=type(e).__name__, ) raise IndexingError( message=f"Failed to delete records for document {document_id}", details={ "document_id": str(document_id), "error": str(e), "error_type": type(e).__name__, }, ) from e