"""Index management and query performance monitoring for DSOL. This module provides the Indexer class for verifying PostgreSQL indexes, monitoring index usage statistics, and benchmarking query performance. """ import time from dataclasses import dataclass from datetime import datetime from typing import Optional from uuid import UUID import structlog from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession logger = structlog.get_logger() @dataclass class IndexInfo: """Information about a single index. Attributes: index_name: Name of the index table_name: Table the index belongs to index_type: Type of index ('gin' or 'btree') column_name: Column(s) the index covers exists: Whether the index exists in the database definition: Full index definition SQL (if exists) """ index_name: str table_name: str index_type: str column_name: str exists: bool definition: Optional[str] = None @dataclass class IndexVerificationReport: """Report of index verification results. Attributes: timestamp: When the verification was performed all_indexes_exist: Whether all expected indexes exist expected_indexes: List of all expected indexes with their status missing_indexes: List of index names that are missing """ timestamp: datetime all_indexes_exist: bool expected_indexes: list[IndexInfo] missing_indexes: list[str] def to_dict(self) -> dict: """Convert to dict for logging. Returns: Dictionary with verification results """ return { "timestamp": self.timestamp.isoformat(), "all_indexes_exist": self.all_indexes_exist, "total_indexes": len(self.expected_indexes), "missing_count": len(self.missing_indexes), "missing_indexes": self.missing_indexes, } @dataclass class IndexStats: """Statistics for a single index. Attributes: index_name: Name of the index scans: Number of times index was used tuples_read: Number of tuples read from index tuples_fetched: Number of tuples fetched from table """ index_name: str scans: int tuples_read: int tuples_fetched: int @dataclass class QueryPerformanceReport: """Query performance benchmark results. Attributes: document_id: Document ID used for testing exact_lookup_ms: Time for exact item code lookup symbol_lookup_ms: Time for symbol lookup jsonb_containment_ms: Time for JSONB containment query resolved_search_ms: Time for resolved content search all_within_targets: Whether all queries meet performance targets """ document_id: str exact_lookup_ms: float symbol_lookup_ms: float jsonb_containment_ms: float resolved_search_ms: float all_within_targets: bool def to_dict(self) -> dict: """Convert to dict for logging. Returns: Dictionary with performance results and target comparisons """ return { "document_id": self.document_id, "exact_lookup_ms": self.exact_lookup_ms, "symbol_lookup_ms": self.symbol_lookup_ms, "jsonb_containment_ms": self.jsonb_containment_ms, "resolved_search_ms": self.resolved_search_ms, "exact_lookup_ok": self.exact_lookup_ms < 50, "symbol_lookup_ok": self.symbol_lookup_ms < 20, "jsonb_containment_ok": self.jsonb_containment_ms < 100, "resolved_search_ok": self.resolved_search_ms < 200, "all_within_targets": self.all_within_targets, } class Indexer: """Manages index verification and query performance monitoring. This class provides methods to: - Verify that expected PostgreSQL indexes exist - Monitor index usage statistics - Benchmark query performance - Verify index usage with EXPLAIN ANALYZE Example: >>> async with get_db_session() as session: ... indexer = Indexer(session) ... report = await indexer.verify_indexes() ... if report.all_indexes_exist: ... print("All indexes OK!") """ # Expected indexes configuration EXPECTED_INDEXES = [ IndexInfo( index_name="idx_extracted_records_content", table_name="extracted_records", index_type="gin", column_name="content", exists=False, ), IndexInfo( index_name="idx_extracted_records_resolved_content", table_name="extracted_records", index_type="gin", column_name="resolved_content", exists=False, ), IndexInfo( index_name="idx_extracted_records_document_id", table_name="extracted_records", index_type="btree", column_name="document_id", exists=False, ), IndexInfo( index_name="idx_extracted_records_sheet_name", table_name="extracted_records", index_type="btree", column_name="sheet_name", exists=False, ), IndexInfo( index_name="idx_symbol_dictionaries_document_id", table_name="symbol_dictionaries", index_type="btree", column_name="document_id", exists=False, ), IndexInfo( index_name="idx_symbol_dictionaries_symbol", table_name="symbol_dictionaries", index_type="btree", column_name="symbol", exists=False, ), ] def __init__(self, db_session: AsyncSession): """Initialize indexer with database session. Args: db_session: Async SQLAlchemy session """ self.db = db_session self.logger = structlog.get_logger() async def verify_indexes(self) -> IndexVerificationReport: """Verify all expected indexes exist. Queries pg_indexes system view to check for: - GIN indexes on extracted_records.content and resolved_content - B-tree indexes on extracted_records.document_id and sheet_name - B-tree indexes on symbol_dictionaries.document_id and symbol Returns: IndexVerificationReport with verification results Example: >>> report = await indexer.verify_indexes() >>> if not report.all_indexes_exist: ... logger.warning("missing_indexes", missing=report.missing_indexes) """ self.logger.info("verifying_indexes", expected_count=len(self.EXPECTED_INDEXES)) # Query pg_indexes for all indexes query = text(""" SELECT indexname, tablename, indexdef FROM pg_indexes WHERE schemaname = 'public' AND tablename IN ('extracted_records', 'symbol_dictionaries') """) result = await self.db.execute(query) existing_indexes = {row[0]: (row[1], row[2]) for row in result.fetchall()} # Check each expected index verified_indexes = [] missing_indexes = [] for expected in self.EXPECTED_INDEXES: if expected.index_name in existing_indexes: table_name, definition = existing_indexes[expected.index_name] verified_indexes.append( IndexInfo( index_name=expected.index_name, table_name=table_name, index_type=expected.index_type, column_name=expected.column_name, exists=True, definition=definition, ) ) else: verified_indexes.append(expected) missing_indexes.append(expected.index_name) all_exist = len(missing_indexes) == 0 report = IndexVerificationReport( timestamp=datetime.now(), all_indexes_exist=all_exist, expected_indexes=verified_indexes, missing_indexes=missing_indexes, ) if all_exist: self.logger.info( "indexes_verified", all_exist=True, total_indexes=len(verified_indexes), ) else: self.logger.warning( "missing_indexes_detected", missing_count=len(missing_indexes), missing_indexes=missing_indexes, ) return report async def get_index_stats( self, table_name: str = "extracted_records" ) -> list[IndexStats]: """Get index usage statistics from PostgreSQL. Queries pg_stat_user_indexes for index usage metrics. Args: table_name: Table name to get stats for (default: extracted_records) Returns: List of IndexStats for each index on the table Example: >>> stats = await indexer.get_index_stats("extracted_records") >>> for stat in stats: ... print(f"{stat.index_name}: {stat.scans} scans") """ self.logger.info("fetching_index_stats", table_name=table_name) query = text(""" SELECT indexrelname AS index_name, idx_scan AS scans, idx_tup_read AS tuples_read, idx_tup_fetch AS tuples_fetched FROM pg_stat_user_indexes WHERE relname = :table_name ORDER BY indexrelname """) result = await self.db.execute(query, {"table_name": table_name}) rows = result.fetchall() stats = [ IndexStats( index_name=row[0], scans=row[1] or 0, tuples_read=row[2] or 0, tuples_fetched=row[3] or 0, ) for row in rows ] self.logger.info( "index_stats_retrieved", table_name=table_name, index_count=len(stats), ) return stats async def benchmark_query_performance( self, document_id: UUID ) -> QueryPerformanceReport: """Benchmark query performance for common patterns. Tests: - Exact item code lookup (target: < 50ms) - Symbol lookup (target: < 20ms) - JSONB containment (target: < 100ms) - Resolved content search (target: < 200ms) Args: document_id: Document ID to use for testing Returns: QueryPerformanceReport with timing results Example: >>> report = await indexer.benchmark_query_performance(doc_id) >>> if not report.all_within_targets: ... logger.warning("performance_below_target", **report.to_dict()) """ self.logger.info("benchmarking_queries", document_id=str(document_id)) # Test 1: Exact item code lookup (target: < 50ms) exact_lookup_ms = await self._benchmark_query( text(""" SELECT * FROM extracted_records WHERE document_id = :document_id AND content->>'Item Code' IS NOT NULL LIMIT 1 """), {"document_id": str(document_id)}, ) # Test 2: Symbol lookup (target: < 20ms) symbol_lookup_ms = await self._benchmark_query( text(""" SELECT * FROM symbol_dictionaries WHERE document_id = :document_id LIMIT 1 """), {"document_id": str(document_id)}, ) # Test 3: JSONB containment (target: < 100ms) jsonb_containment_ms = await self._benchmark_query( text(""" SELECT * FROM extracted_records WHERE document_id = :document_id AND content IS NOT NULL LIMIT 10 """), {"document_id": str(document_id)}, ) # Test 4: Resolved content search (target: < 200ms) resolved_search_ms = await self._benchmark_query( text(""" SELECT * FROM extracted_records WHERE document_id = :document_id AND resolved_content IS NOT NULL LIMIT 10 """), {"document_id": str(document_id)}, ) # Check if all within targets all_within_targets = ( exact_lookup_ms < 50 and symbol_lookup_ms < 20 and jsonb_containment_ms < 100 and resolved_search_ms < 200 ) report = QueryPerformanceReport( document_id=str(document_id), exact_lookup_ms=exact_lookup_ms, symbol_lookup_ms=symbol_lookup_ms, jsonb_containment_ms=jsonb_containment_ms, resolved_search_ms=resolved_search_ms, all_within_targets=all_within_targets, ) self.logger.info("benchmark_complete", **report.to_dict()) return report async def _benchmark_query(self, query: text, params: dict) -> float: """Execute a query and measure its execution time. Args: query: SQL query to execute params: Query parameters Returns: Execution time in milliseconds """ start_time = time.perf_counter() await self.db.execute(query, params) end_time = time.perf_counter() duration_ms = (end_time - start_time) * 1000 return round(duration_ms, 2) async def explain_query(self, query_sql: str, params: Optional[dict] = None) -> dict: """Run EXPLAIN ANALYZE on a query to verify index usage. Args: query_sql: SQL query to analyze params: Optional query parameters Returns: Dict with query plan and execution statistics Example: >>> plan = await indexer.explain_query( ... "SELECT * FROM extracted_records WHERE content->>'Item Code' = :code", ... {"code": "2024-4_0019"} ... ) >>> if "Index Scan" in plan["plan_text"]: ... print("Using index!") """ self.logger.info("explaining_query", query_preview=query_sql[:100]) # Wrap query in EXPLAIN ANALYZE explain_query = text(f"EXPLAIN ANALYZE {query_sql}") result = await self.db.execute(explain_query, params or {}) plan_rows = result.fetchall() # Combine plan rows into single text plan_text = "\n".join(row[0] for row in plan_rows) # Extract key information uses_index = "Index Scan" in plan_text or "Bitmap Index Scan" in plan_text uses_sequential_scan = "Seq Scan" in plan_text plan_info = { "plan_text": plan_text, "uses_index": uses_index, "uses_sequential_scan": uses_sequential_scan, "query_preview": query_sql[:200], } self.logger.info( "query_explained", uses_index=uses_index, uses_sequential_scan=uses_sequential_scan, ) return plan_info