添加 羽心堂只筛选指定仓库

inventory_sync/postgres.py

@@ -7,10 +7,26 @@ import subprocess
 from typing import Any
 
 from .settings import AppConfig
-from .sql import POSTGRES_QUERY
+from .specs import BrandSyncTarget
+from .sql import DEFAULT_SYNC_TARGET, POSTGRES_QUERY, build_postgres_query
 
 
-def query_postgres(config: AppConfig) -> dict[str, Any]:
+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 = [
@@ -28,7 +44,7 @@ def query_postgres(config: AppConfig) -> dict[str, Any]:
         "-c",
         "SET TIME ZONE 'Asia/Shanghai';",
         "-c",
-        f"COPY ({POSTGRES_QUERY}) TO STDOUT WITH CSV HEADER",
+        f"COPY ({query}) TO STDOUT WITH CSV HEADER",
     ]
 
     try:
@@ -48,4 +64,3 @@ def query_postgres(config: AppConfig) -> dict[str, Any]:
     )
     rows = list(csv.DictReader(io.StringIO(cleaned_output)))
     return {"row_count": len(rows), "rows": rows}
-