"""Results batcher.""" from typing import Any, cast from abacus_common_logic.concurrent import ThreadSafeIterableIterator from src.connectors.db import ColumnDescriptionIndex, Connection, Cursor class DBResultsIterable[T = Any](ThreadSafeIterableIterator[list[T]]): """Batched database query results iterator. Executes a query and yields results in chunks (batches) of a specified size. Extends ThreadSafeIterableIterator to allow safe concurrent consumption of batches by multiple worker threads. Lazy execution: The query is not executed until the first batch is requested. """ def __init__(self, conn: Connection, query: str, batch_size: int): """Initialize the results batcher. Args: conn: Active database connection. query: SQL query to execute. batch_size: Number of rows to yield per batch. Raises: ValueError: If batch_size is not positive. """ super().__init__() if batch_size < 1: raise ValueError(f'batch_size must be positive. Received: {batch_size}') self._batch_size = batch_size self._conn = conn self._cursor: Cursor | None = None self._query = query def _next(self) -> list[T]: """Fetch the next batch of rows. Returns: list[T]: A list of rows (size <= batch_size). Raises: StopIteration: If there are no more rows. """ batch = cast(Cursor, self._cursor).fetchmany(self._batch_size) if not batch: raise StopIteration return cast(list[T], batch) def _start(self) -> None: """Execute the query (lazy initialization).""" if self._cursor is None: self._cursor = cast(Cursor, self._conn.cursor()) self._cursor.execute(self._query) def _destroy(self) -> None: """Close the cursor when iteration finishes or fails.""" if not self._cursor: return try: self._cursor.close() except Exception: pass self._cursor = None def calculate_batch_size( bytes_per_entry: int, base_query_length: int, max_query_length: int, safety_margin_pct: float = 1.0, max_batch_size: int | None = None, ) -> int: """Calculate optimal batch size based on packet limits and entry size. The calculation reserves space for the base query and applies a safety margin. It accounts for 1 byte of separator overhead (comma) per entry. Args: bytes_per_entry: Estimated byte size of a single data entry. base_query_length: Byte length of the SQL query (e.g., "INSERT INTO..."). max_query_length: Max query length in bytes. safety_margin_pct: Fraction of packet limit to utilize (0.0 to 1.0). Defaults to DBConfig.BATCH_SAFETY_MARGIN_PCT. max_batch_size: Optional max batch size. Returns: int: Optimal batch size. Raises: ValueError: If inputs are invalid. """ if bytes_per_entry < 1: raise ValueError( f'bytes_per_entry must be positive. Received: {bytes_per_entry}' ) if base_query_length < 0: raise ValueError( f'base_query_length must be nonnegative. Received: {base_query_length}' ) if not 0 <= safety_margin_pct <= 1: raise ValueError( f'safety_margin_pct must be within [0, 1]. Received: {safety_margin_pct}' ) if max_query_length < 1: raise ValueError( f'max_query_length must be positive. Received: {max_query_length}' ) # Calculate safe available bytes avail_bytes = int(max_query_length * safety_margin_pct) - base_query_length # Calculate batch size # Adds 1 byte to bytes_per_entry for the comma separator: 1,2,3,... # However, the last item has no comma: ...7,8,9 # 1 byte is added to avail_bytes to adjust for the additional comma. batch_size = int((avail_bytes + 1) / (bytes_per_entry + 1)) # Apply max bound if max_batch_size: batch_size = min(max_batch_size, batch_size) # Apply min bound return max(0, batch_size) def count_query_rows(cursor: Cursor, query: str) -> int: """Get the number of rows in a query result. Args: cursor: Active database connection cursor. query: The query. Returns: int: The number of rows. """ result = cursor.execute(f'SELECT COUNT(*) FROM ({query}) AS t').fetchone() return int(result[0]) if result else 0 def count_table_rows(cursor: Cursor, table_name: str) -> int: """Get the number of rows in a database table. Args: cursor: Active database connection cursor. table_name: The name of the table. Returns: int: The number of rows. """ result = cursor.execute(f'SELECT COUNT(*) FROM "{table_name}"').fetchone() return int(result[0]) if result else 0 def drop_table(cursor: Cursor, table_name: str) -> None: """Drop a table or view if it exists. DuckDB requires specific commands for Tables vs Views. This attempts to drop both to ensure the name is dropped. Args: cursor: Active database connection cursor. table_name: The name of the object to drop. """ cursor.execute(f'DROP TABLE IF EXISTS "{table_name}"') cursor.execute(f'DROP VIEW IF EXISTS "{table_name}"') def get_result_dict(cursor: Cursor) -> dict[str, Any]: """Fetch one row and converts it into a dictionary. Args: cursor: The executed database cursor. Returns: A dictionary of the result row, or an empty dictionary if no row is found. """ result = cursor.fetchone() if result is None or cursor.description is None: return {} column_names: list[str] = [] for desc in cursor.description: column_names.append(cast(str, desc[ColumnDescriptionIndex.NAME])) return dict(zip(column_names, result)) def get_select_clause(alias_map: dict[str, str | None] | None = None) -> str: """Build SELECT clause for column mapping. Args: alias_map: Optional mapping from alias to source column. Returns: str: SELECT clause; either '*' or 'col1 AS new1, col2 AS new2, ...'. """ if not alias_map: return '*' parts: list[str] = [] for alias, src in alias_map.items(): src = f'"{src}"' if src else 'CAST(NULL AS VARCHAR)' parts.append(f'{src} AS "{alias}"') return ', '.join(parts) def merge_tables(cursor: Cursor, table_name: str, sub_tables: list[str]) -> int: """Merge multiple temporary tables into a single table. Creates a new table from the first sub-table, then inserts data from remaining sub-tables sequentially. Drops each sub-table after merging. Args: cursor: Active database connection cursor. table_name: Name for the merged table. sub_tables: List of temporary table names to merge. Returns: Total row count in the merged table. """ if not sub_tables: return 0 cursor.execute( f""" CREATE TABLE "{table_name}" AS SELECT * FROM "{sub_tables[0]}" """ ) cursor.execute(f'DROP TABLE "{sub_tables[0]}"') for table in sub_tables[1:]: cursor.execute( f""" INSERT INTO "{table_name}" SELECT * FROM "{table}" """ ) cursor.execute(f'DROP TABLE "{table}"') return count_table_rows(cursor, table_name)