修复 prefect中psql报错

inventory_sync/postgres.py

@@ -2,10 +2,10 @@ from __future__ import annotations
 
 import csv
 import io
-import os
-import subprocess
 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
@@ -27,40 +27,21 @@ def query_postgres(
             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",
-    ]
-
+    output = io.StringIO()
     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
+        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
 
-    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)))
+    output.seek(0)
+    rows = list(csv.DictReader(output))
     return {"row_count": len(rows), "rows": rows}