"""Data source module.""" import tempfile from enum import Enum, auto from pathlib import Path from typing import Any, NamedTuple from abacus_common_logic.concurrent import RelayRunner from duckdb import DuckDBPyConnection as DuckDBPyCursor from config import config from src.connectors.duckdb import DuckDBConnection from src.connectors.duckdb.utils import create_parquet_from_query from src.connectors.snowflake import ( SnowflakeConnection, get_max_query_bytes, get_max_query_values, ) from src.enums import DuckDBTable, Environment from src.gateways.sql import SnowflakeQuery, load_sql from src.infra.log import logger from src.infra.resources import ResourceManager from src.sql import DuckDBQuery, load_sql as load_duckdb_sql from src.utils.db_utils import ( DBResultsIterable, calculate_batch_size, count_query_rows, count_table_rows, drop_table, merge_tables, ) from src.utils.file_utils import gen_id class LoadStrategy(Enum): """Loading strategy for reference tables.""" BULK = auto() BATCH = auto() EMPTY = auto() class BatchConfig(NamedTuple): """Batch loading configuration with query templates and ID sizing.""" bytes_per_id: int batch_query: str id_query: str count_query: str | None = None placeholder: str = '{values}' class LoadConfig(NamedTuple): """Table loading configuration with source query and optional batching.""" query: str batch: BatchConfig | None = None class SnowflakeGateway: """Snowflake data source implementation.""" def __init__( self, local_conn: DuckDBConnection, remote_conn: SnowflakeConnection, secret_name: str, env: Environment, ): """Initialize the source.""" self._duck_conn = local_conn self._env = env if env.is_managed else Environment.QA self._ref_tables = self._init_ref_tables() self._secret = secret_name self._snow_conn = remote_conn def import_accounts(self) -> int: """Import accounts.""" return self._import(DuckDBTable.ACCOUNT) def import_account_upcs(self) -> int: """Import account UPCs.""" return self._import(DuckDBTable.ACCOUNT_UPCS) def import_account_contracts(self) -> int: """Import account contracts.""" return self._import(DuckDBTable.ACCOUNT_CONTRACT) def import_account_payment_terms(self) -> int: """Import account payment terms.""" return self._import(DuckDBTable.ACCOUNT_PAYMENT_TERM) def import_close_balance_statuses(self) -> int: """Import close balance statuses.""" return self._import(DuckDBTable.CLOSE_BALANCE_STATUS) def import_currency_codes(self) -> int: """Import currency codes.""" return self._import(DuckDBTable.CURRENCY_CODE) def import_flat_contract_terms(self) -> int: """Import flattened contract terms.""" return self._import(DuckDBTable.FLAT_CONTRACT_TERM) def import_reference_adjustment_types(self) -> int: """Import reference adjustment types.""" return self._import(DuckDBTable.REF_ADJUSTMENT_TYPE) def import_statement_periods(self) -> int: """Import statement periods.""" return self._import(DuckDBTable.STATEMENT_PERIOD) def _create_from_remote( self, cursor: DuckDBPyCursor, table_name: str, query: str ) -> int: """Create from Snowflake. Args: cursor: Active DuckDB connection. query: The Snowflake query to source data from. table_name: The name of the table to create. Returns: int: The number of rows created. """ query = query.replace('{env}', self._env) outer_query = f""" CREATE OR REPLACE TABLE "{table_name}" AS SELECT * FROM snowflake_query($query, $secret) """ cursor.execute(outer_query, {'query': query, 'secret': self._secret}) return count_table_rows(cursor, table_name) def _export_to_remote(self, table_name: str, file_path: str | Path) -> int: rid = gen_id() stage_name = f'tmp_stage_{rid}' file_format = f'tmp_format_{rid}' with self._snow_conn.cursor() as cursor: try: logger.info(f'Creating temporary Snowflake stage: {stage_name}') cursor.execute(f'CREATE OR REPLACE TEMPORARY STAGE {stage_name}') cursor.execute(f""" CREATE OR REPLACE TEMPORARY FILE FORMAT {file_format} TYPE = PARQUET COMPRESSION = AUTO USE_LOGICAL_TYPE = TRUE """) logger.info(f'Uploading file to Snowflake: {file_path}') cursor.execute(f""" PUT 'file://{file_path}' @{stage_name} AUTO_COMPRESS=TRUE """) logger.info(f'Creating Snowflake table: {table_name}') cursor.execute(f""" CREATE OR REPLACE TRANSIENT TABLE {table_name} USING TEMPLATE ( SELECT ARRAY_AGG(OBJECT_CONSTRUCT(*)) FROM TABLE( INFER_SCHEMA( LOCATION=>'@{stage_name}', FILE_FORMAT=>'{file_format}', IGNORE_CASE=>TRUE ) ) ) DATA_RETENTION_TIME_IN_DAYS = 0 """) logger.info('Populating Snowflake table from stage') cursor.execute(f""" COPY INTO {table_name} FROM @{stage_name} FILE_FORMAT = ( FORMAT_NAME = '{file_format}' ) MATCH_BY_COLUMN_NAME = CASE_INSENSITIVE PURGE = TRUE """) cursor.execute(f'SELECT COUNT(*) FROM {table_name}') result = cursor.fetchone() row_count = result[0] if result else 0 logger.info(f'Uploaded {table_name}: {row_count} rows') return row_count finally: cursor.execute(f'DROP STAGE IF EXISTS {stage_name}') cursor.execute(f'DROP FILE FORMAT IF EXISTS {file_format}') def _get_import_strategy( self, table_name: DuckDBTable, config: LoadConfig ) -> LoadStrategy: # No batching configured if not config.batch: return LoadStrategy.BULK # Count total rows with self._duck_conn.cursor() as cursor: count_query = config.batch.count_query if not count_query: count_query = f'SELECT COUNT(*) FROM ({config.query})' result = self._query_remote_one(cursor, count_query) total_count = result[0] if result else 0 logger.info(f'{table_name}: {total_count} total remote rows') # Strategy 1: Small table if total_count <= self._get_max_bulk_rows(): return LoadStrategy.BULK # Count local unique IDs id_query = config.batch.id_query with self._duck_conn.cursor() as cursor: id_count = count_query_rows(cursor, id_query) pct_total = float(id_count / total_count if total_count > 0 else 0) logger.info(f'{table_name}: {id_count} local IDs ({(100.0 * pct_total):.2f}%)') # Strategy 2: No local IDs if id_count < 1: return LoadStrategy.EMPTY # Strategy 3: Filtering is efficient if pct_total < self._get_min_bulk_ratio(): return LoadStrategy.BATCH # Strategy 4: Filtering is not efficient return LoadStrategy.BULK def _get_max_bulk_rows(self) -> int: """Get maximum rows for bulk loading.""" return config.exec.MAX_BULK_ROWS def _get_max_query_filter_ratio(self) -> float: """Get maximum ratio for query filtering.""" return config.exec.MAX_QUERY_FILTER_RATIO def _get_min_bulk_ratio(self) -> float: """Get minimum ratio for bulk loading.""" return config.exec.MIN_BULK_RATIO def _init_ref_tables(self) -> dict[DuckDBTable, LoadConfig]: """Initialize reference table configurations. Returns: Mapping of table names to their loading configurations. """ return { DuckDBTable.ACCOUNT: LoadConfig( query=load_sql(SnowflakeQuery.GetAccounts), ), DuckDBTable.ACCOUNT_CONTRACT: LoadConfig( query=load_sql(SnowflakeQuery.GetAccountContracts), ), DuckDBTable.ACCOUNT_PAYMENT_TERM: LoadConfig( query=load_sql(SnowflakeQuery.GetAccountPaymentTerms), ), DuckDBTable.ACCOUNT_UPCS: LoadConfig( query=load_sql(SnowflakeQuery.BulkAccountUpcs), batch=BatchConfig( # 14 bytes per upc + 2 escaped quotes bytes_per_id=18, batch_query=load_sql(SnowflakeQuery.BatchAccountUpcs), count_query=load_sql(SnowflakeQuery.CountAccountUpcs), id_query=load_duckdb_sql(DuckDBQuery.GetUniqueUpcs), ), ), DuckDBTable.CLOSE_BALANCE_STATUS: LoadConfig( query=load_sql(SnowflakeQuery.GetCloseBalanceStatuses), ), DuckDBTable.CURRENCY_CODE: LoadConfig( query=load_sql(SnowflakeQuery.GetCurrencyCodes), ), DuckDBTable.FLAT_CONTRACT_TERM: LoadConfig( query=load_sql(SnowflakeQuery.GetFlatContractTerms), ), DuckDBTable.REF_ADJUSTMENT_TYPE: LoadConfig( query=load_sql(SnowflakeQuery.GetReferenceAdjustmentTypes), ), DuckDBTable.STATEMENT_PERIOD: LoadConfig( query=load_sql(SnowflakeQuery.GetStatementPeriods), ), } def _import( self, table_name: DuckDBTable, ) -> int: """Load a single reference table based on configuration. Args: table_name: The name of the table to create in DuckDB. config: The configuration defining the source query and batching rules. executor: Executor for concurrent batch operations. Returns: Number of rows loaded. """ config = self._ref_tables[table_name] strategy = self._get_import_strategy(table_name, config) if strategy == LoadStrategy.BULK: return self._import_bulk(table_name, config.query) if strategy == LoadStrategy.BATCH: return self._import_staged(table_name, config.batch) # type: ignore # return self._batch_load(table_name, config.batch, executor) # type: ignore if strategy == LoadStrategy.EMPTY: limited_query = f'SELECT * FROM ({config.query}) LIMIT 0' return self._import_bulk(table_name, limited_query) raise ValueError(f"Unknown strategy '{strategy}'") def _import_batched(self, table_name: DuckDBTable, batch: BatchConfig) -> int: """Parallel download into temp tables with a sliding window, then sequential merge. Args: table_name: The name of the table to create in DuckDB. batch: The configuration defining batching rules. Returns: int: Total number of rows loaded. """ batch_size = calculate_batch_size( bytes_per_entry=batch.bytes_per_id, base_query_length=len(batch.batch_query), max_query_length=get_max_query_bytes(), safety_margin_pct=self._get_max_query_filter_ratio(), max_batch_size=get_max_query_values(), ) if batch_size < 1: raise ValueError( f'Calculated batch size is {batch_size}. ' f'Query length ({len(batch.batch_query)}) is too large relative to ' f'max query length and entry size ({batch.bytes_per_id}).' ) logger.info(f'Batch loading {table_name}. Batch size: {batch_size}') id_batches = DBResultsIterable(self._duck_conn, batch.id_query, batch_size) batch_tables: list[str] = [] placeholder = batch.placeholder if batch.placeholder else '{values}' query = batch.batch_query def _load_batch(batch: list[Any], batch_index: int) -> None: tmp_table_name = f'tmp_{table_name}_{batch_index}' # Append is thread-safe, and there is no surrounding logic batch_tables.append(tmp_table_name) snowflake_sql = self._populate_batch_query(query, batch, placeholder) logger.info(f'Loading batch {tmp_table_name}') with self._duck_conn.cursor() as cursor: row_count = self._create_from_remote( cursor, tmp_table_name, snowflake_sql ) logger.info(f'Loaded batch {tmp_table_name}: {row_count} rows') executor = ResourceManager.get_io_pool() relay = RelayRunner(executor, id_batches, _load_batch) relay.run() logger.info(f'Merging {len(batch_tables)} batches for {table_name}') with self._duck_conn.cursor() as cursor: return merge_tables(cursor, table_name, batch_tables) def _import_bulk(self, table_name: DuckDBTable, query: str) -> int: """Perform bulk load of table from Snowflake.""" logger.info(f'Bulk loading {table_name}') with self._duck_conn.cursor() as cursor: row_count = self._create_from_remote(cursor, table_name, query) logger.info(f'Loaded {table_name}: {row_count} rows') return row_count def _import_staged(self, table_name: DuckDBTable, batch: BatchConfig) -> int: """Create table from a parquet file.""" logger.info(f'Remote loading {table_name}') tmp = config.storage.TEMP_DIR with tempfile.NamedTemporaryFile(dir=tmp, suffix='.parquet') as fd: fp = Path(fd.name) with self._duck_conn.cursor() as cursor: logger.info(f'Exporting IDs to file: {fp}') create_parquet_from_query(cursor, batch.id_query, fp, True) temp_table = f'DEV_ENGINEERING.SCRATCH.tmp_{table_name}_{gen_id()}' try: row_count = self._export_to_remote(temp_table, fp) logger.info(f'Loading to: {table_name}') batch_query = batch.batch_query.replace('{temp_table}', temp_table) row_count = self._create_from_remote( cursor, table_name, batch_query ) logger.info(f'Loaded {table_name}: {row_count} rows') finally: logger.info(f'Dropping table: {temp_table}') drop_table(cursor, temp_table) logger.info(f'Dropped table: {temp_table}') return row_count def _populate_batch_query( self, query: str, batch: list[Any], placeholder: str = '{values}' ) -> str: """Create a query for fetching a batch of data from Snowflake. Args: query: The SQL query template with a placeholder. batch: A list of values to substitute into the query. placeholder: The placeholder string to replace. Returns: str: The query with the placeholder replaced by the comma-separated values. """ if batch and isinstance(batch[0], (list, tuple)): batch = [row[0] for row in batch] escaped_values = [f"('{str(val).replace("'", "''")}')" for val in batch] values_str = ','.join(escaped_values) return query.replace(placeholder, values_str) def _query_remote_one( self, cursor: DuckDBPyCursor, query: str ) -> tuple[Any, ...] | None: """Execute a query in Snowflake and return the first result. Args: cursor: Active DuckDB connection cursor. query: The Snowflake query to execute. Returns: int: The first value of the first row of the result. """ query = query.replace('{env}', self._env) outer_query = 'SELECT * FROM snowflake_query($query, $secret)' cursor.execute(outer_query, {'query': query, 'secret': self._secret}) return cursor.fetchone()