from __future__ import annotations import csv import io import logging import time from typing import Any import psycopg2 from inventory_analytics.logging_utils import log_event as emit_log_event from .settings import AppConfig from .specs import BrandSyncTarget from .sql import DEFAULT_SYNC_TARGET, POSTGRES_QUERY, build_postgres_query logger = logging.getLogger("inventory_sync.postgres") LOG_FIELD_LABELS = { "elapsed_ms": "耗时毫秒", "row_count": "行数", "stage": "阶段", "statement_timeout_ms": "查询超时毫秒", } def log_event(level: int, event: str, **fields: Any) -> None: emit_log_event(logger, LOG_FIELD_LABELS, level, event, **fields) def query_postgres( config: AppConfig, brand_name: str | None = None, warehouse_keys: tuple[int | str, ...] | None = None, target: BrandSyncTarget | None = None, ) -> dict[str, Any]: started_at = time.monotonic() statement_timeout_ms = 300_000 active_brand_name = brand_name or (target.brand_name if target else DEFAULT_SYNC_TARGET.brand_name) active_warehouse_keys = warehouse_keys or (target.warehouse_keys if target else DEFAULT_SYNC_TARGET.warehouse_keys) query = ( POSTGRES_QUERY if brand_name is None and warehouse_keys is None and target is None else build_postgres_query( brand_name=active_brand_name, warehouse_keys=active_warehouse_keys, ) ) output = io.StringIO() try: log_event(logging.INFO, "开始连接PostgreSQL", stage="database_connect") with psycopg2.connect( host=config.analytics_db_host, port=config.analytics_db_port, user=config.analytics_db_user, password=config.analytics_db_password, dbname=config.analytics_db_name, connect_timeout=10, ) as conn: log_event( logging.INFO, "PostgreSQL连接完成", stage="database_connect", elapsed_ms=int((time.monotonic() - started_at) * 1000), ) with conn.cursor() as cur: cur.execute("SET TIME ZONE 'Asia/Shanghai';") cur.execute("SET statement_timeout = %s;", (statement_timeout_ms,)) copy_started_at = time.monotonic() log_event( logging.INFO, "开始执行渠道库存SQL", stage="database_query", statement_timeout_ms=statement_timeout_ms, ) cur.copy_expert(f"COPY ({query}) TO STDOUT WITH CSV HEADER", output) log_event( logging.INFO, "渠道库存SQL执行完成", stage="database_query", elapsed_ms=int((time.monotonic() - copy_started_at) * 1000), ) except psycopg2.Error as exc: raise RuntimeError(f"PostgreSQL 查询失败: {exc}") from exc parse_started_at = time.monotonic() output.seek(0) rows = list(csv.DictReader(output)) log_event( logging.INFO, "渠道库存查询结果解析完成", stage="database_parse", row_count=len(rows), elapsed_ms=int((time.monotonic() - parse_started_at) * 1000), ) return {"row_count": len(rows), "rows": rows}