37 lines
1.5 KiB
Python
37 lines
1.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""初始化 DWD Schema：执行 schema_dwd_doc.sql，可选先 DROP SCHEMA。"""
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from .base_task import BaseTask, TaskContext
|
||
|
||
|
||
class InitDwdSchemaTask(BaseTask):
|
||
"""通过调度执行 DWD schema 初始化。"""
|
||
|
||
def get_task_code(self) -> str:
|
||
"""返回任务编码。"""
|
||
return "INIT_DWD_SCHEMA"
|
||
|
||
def extract(self, context: TaskContext) -> dict[str, Any]:
|
||
"""读取 DWD SQL 文件与参数。"""
|
||
base_dir = Path(__file__).resolve().parents[1] / "database"
|
||
dwd_path = Path(self.config.get("schema.dwd_file", base_dir / "schema_dwd_doc.sql"))
|
||
if not dwd_path.exists():
|
||
raise FileNotFoundError(f"未找到 DWD schema 文件: {dwd_path}")
|
||
|
||
drop_first = self.config.get("dwd.drop_schema_first", False)
|
||
return {"dwd_sql": dwd_path.read_text(encoding="utf-8"), "dwd_file": str(dwd_path), "drop_first": drop_first}
|
||
|
||
def load(self, extracted: dict[str, Any], context: TaskContext) -> dict:
|
||
"""可选 DROP schema，再执行 DWD DDL。"""
|
||
with self.db.conn.cursor() as cur:
|
||
if extracted["drop_first"]:
|
||
cur.execute("DROP SCHEMA IF EXISTS billiards_dwd CASCADE;")
|
||
self.logger.info("已执行 DROP SCHEMA billiards_dwd CASCADE")
|
||
self.logger.info("执行 DWD schema 文件: %s", extracted["dwd_file"])
|
||
cur.execute(extracted["dwd_sql"])
|
||
return {"executed": 1, "files": [extracted["dwd_file"]]}
|