inventory_analytics/inventory_sync/postgres.py

from __future__ import annotations

import csv
import io
import os
import subprocess
from typing import Any

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,
        )
    )
    env = os.environ.copy()
    env["PGPASSWORD"] = config.analytics_db_password
    command = [
        "/usr/local/bin/psql",
        "-h",
        config.analytics_db_host,
        "-p",
        config.analytics_db_port,
        "-U",
        config.analytics_db_user,
        "-d",
        config.analytics_db_name,
        "-v",
        "ON_ERROR_STOP=1",
        "-c",
        "SET TIME ZONE 'Asia/Shanghai';",
        "-c",
        f"COPY ({query}) TO STDOUT WITH CSV HEADER",
    ]

    try:
        completed = subprocess.run(
            command,
            check=True,
            capture_output=True,
            text=True,
            env=env,
        )
    except subprocess.CalledProcessError as exc:
        detail = exc.stderr.strip() or exc.stdout.strip()
        raise RuntimeError(f"PostgreSQL 查询失败: {detail}") from exc

    cleaned_output = "\n".join(
        line for line in completed.stdout.splitlines() if line.strip() and line.strip() != "SET"
    )
    rows = list(csv.DictReader(io.StringIO(cleaned_output)))
    return {"row_count": len(rows), "rows": rows}