inventory_analytics/inventory_sync/postgres.py

from __future__ import annotations

import csv
import io
from typing import Any

import psycopg2

from .settings import AppConfig
from .specs import BrandSyncTarget
from .sql import DEFAULT_SYNC_TARGET, POSTGRES_QUERY, build_postgres_query


def query_postgres(
    config: AppConfig,
    brand_name: str | None = None,
    warehouse_keys: tuple[int | str, ...] | None = None,
    target: BrandSyncTarget | None = None,
) -> dict[str, Any]:
    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:
        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,
        ) as conn:
            with conn.cursor() as cur:
                cur.execute("SET TIME ZONE 'Asia/Shanghai';")
                cur.copy_expert(f"COPY ({query}) TO STDOUT WITH CSV HEADER", output)
    except psycopg2.Error as exc:
        raise RuntimeError(f"PostgreSQL 查询失败: {exc}") from exc

    output.seek(0)
    rows = list(csv.DictReader(output))
    return {"row_count": len(rows), "rows": rows}