67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
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}
|