inventory_analytics/inventory_sync/postgres.py

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}