48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
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}
|