合并

etl_billiards/tasks/init_dws_schema_task.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+"""Initialize DWS schema (billiards_dws)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from .base_task import BaseTask, TaskContext
+
+
+class InitDwsSchemaTask(BaseTask):
+    """Apply DWS schema SQL."""
+
+    def get_task_code(self) -> str:
+        return "INIT_DWS_SCHEMA"
+
+    def extract(self, context: TaskContext) -> dict[str, Any]:
+        base_dir = Path(__file__).resolve().parents[1] / "database"
+        dws_path = Path(self.config.get("schema.dws_file", base_dir / "schema_dws.sql"))
+        if not dws_path.exists():
+            raise FileNotFoundError(f"未找到 DWS schema 文件: {dws_path}")
+        drop_first = bool(self.config.get("dws.drop_schema_first", False))
+        return {"dws_sql": dws_path.read_text(encoding="utf-8"), "dws_file": str(dws_path), "drop_first": drop_first}
+
+    def load(self, extracted: dict[str, Any], context: TaskContext) -> dict:
+        with self.db.conn.cursor() as cur:
+            if extracted["drop_first"]:
+                cur.execute("DROP SCHEMA IF EXISTS billiards_dws CASCADE;")
+                self.logger.info("已执行 DROP SCHEMA billiards_dws CASCADE")
+            self.logger.info("执行 DWS schema 文件: %s", extracted["dws_file"])
+            cur.execute(extracted["dws_sql"])
+        return {"executed": 1, "files": [extracted["dws_file"]]}
+