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}