98 lines
3.3 KiB
Python
98 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import csv
|
|
import io
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
import psycopg2
|
|
|
|
from inventory_analytics.logging_utils import log_event as emit_log_event
|
|
|
|
from .settings import AppConfig
|
|
from .specs import BrandSyncTarget
|
|
from .sql import DEFAULT_SYNC_TARGET, POSTGRES_QUERY, build_postgres_query
|
|
|
|
|
|
logger = logging.getLogger("inventory_sync.postgres")
|
|
|
|
LOG_FIELD_LABELS = {
|
|
"elapsed_ms": "耗时毫秒",
|
|
"row_count": "行数",
|
|
"stage": "阶段",
|
|
"statement_timeout_ms": "查询超时毫秒",
|
|
}
|
|
|
|
|
|
def log_event(level: int, event: str, **fields: Any) -> None:
|
|
emit_log_event(logger, LOG_FIELD_LABELS, level, event, **fields)
|
|
|
|
|
|
def query_postgres(
|
|
config: AppConfig,
|
|
brand_name: str | None = None,
|
|
warehouse_keys: tuple[int | str, ...] | None = None,
|
|
target: BrandSyncTarget | None = None,
|
|
) -> dict[str, Any]:
|
|
started_at = time.monotonic()
|
|
statement_timeout_ms = 300_000
|
|
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:
|
|
log_event(logging.INFO, "开始连接PostgreSQL", stage="database_connect")
|
|
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,
|
|
connect_timeout=10,
|
|
) as conn:
|
|
log_event(
|
|
logging.INFO,
|
|
"PostgreSQL连接完成",
|
|
stage="database_connect",
|
|
elapsed_ms=int((time.monotonic() - started_at) * 1000),
|
|
)
|
|
with conn.cursor() as cur:
|
|
cur.execute("SET TIME ZONE 'Asia/Shanghai';")
|
|
cur.execute("SET statement_timeout = %s;", (statement_timeout_ms,))
|
|
copy_started_at = time.monotonic()
|
|
log_event(
|
|
logging.INFO,
|
|
"开始执行渠道库存SQL",
|
|
stage="database_query",
|
|
statement_timeout_ms=statement_timeout_ms,
|
|
)
|
|
cur.copy_expert(f"COPY ({query}) TO STDOUT WITH CSV HEADER", output)
|
|
log_event(
|
|
logging.INFO,
|
|
"渠道库存SQL执行完成",
|
|
stage="database_query",
|
|
elapsed_ms=int((time.monotonic() - copy_started_at) * 1000),
|
|
)
|
|
except psycopg2.Error as exc:
|
|
raise RuntimeError(f"PostgreSQL 查询失败: {exc}") from exc
|
|
|
|
parse_started_at = time.monotonic()
|
|
output.seek(0)
|
|
rows = list(csv.DictReader(output))
|
|
log_event(
|
|
logging.INFO,
|
|
"渠道库存查询结果解析完成",
|
|
stage="database_parse",
|
|
row_count=len(rows),
|
|
elapsed_ms=int((time.monotonic() - parse_started_at) * 1000),
|
|
)
|
|
return {"row_count": len(rows), "rows": rows}
|