376 lines
14 KiB
Python
376 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
DWS 历史数据重算脚本 — 按营业日（Business Day Cutoff）口径重新聚合
|
||
|
||
用途：
|
||
引入 BUSINESS_DAY_START_HOUR 后，需要将历史 DWS 数据按新口径重算。
|
||
脚本复用正式 ETL CLI（cli.main），按月分窗口执行，避免单次事务过大。
|
||
|
||
用法：
|
||
python scripts/ops/rebuild_dws_biz_date.py --start-date 2025-11-01 --end-date 2026-03-01
|
||
python scripts/ops/rebuild_dws_biz_date.py --start-date 2025-11-01 --end-date 2026-03-01 --dry-run
|
||
python scripts/ops/rebuild_dws_biz_date.py --start-date 2025-11-01 --end-date 2026-03-01 --fail-fast
|
||
python scripts/ops/rebuild_dws_biz_date.py --start-date 2025-11-01 --end-date 2026-03-01 --tasks DWS_FINANCE_DAILY,DWS_MEMBER_VISIT
|
||
|
||
Requirements: 10.1, 10.2, 10.3, 10.4
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
from datetime import date, datetime, timedelta
|
||
from pathlib import Path
|
||
from zoneinfo import ZoneInfo
|
||
|
||
from dotenv import load_dotenv
|
||
|
||
# ── 环境初始化 ────────────────────────────────────────────────
|
||
_ROOT = Path(__file__).resolve().parents[2]
|
||
load_dotenv(_ROOT / ".env")
|
||
|
||
TZ = ZoneInfo("Asia/Shanghai")
|
||
ETL_CWD = _ROOT / "apps" / "etl" / "connectors" / "feiqiu"
|
||
|
||
# 所有需要重算的 DWS 任务代码（按依赖顺序排列）
|
||
# 不含 DWS_MAINTENANCE / INDEX 层——重算完成后可单独刷新
|
||
ALL_DWS_TASKS: list[str] = [
|
||
# 财务维度
|
||
"DWS_FINANCE_DAILY",
|
||
"DWS_FINANCE_RECHARGE",
|
||
"DWS_FINANCE_INCOME_STRUCTURE",
|
||
"DWS_FINANCE_DISCOUNT_DETAIL",
|
||
# 助教维度
|
||
"DWS_ASSISTANT_DAILY",
|
||
"DWS_ASSISTANT_ORDER_CONTRIBUTION",
|
||
"DWS_ASSISTANT_CUSTOMER",
|
||
"DWS_ASSISTANT_SALARY",
|
||
"DWS_ASSISTANT_FINANCE",
|
||
"DWS_ASSISTANT_MONTHLY",
|
||
# 会员维度
|
||
"DWS_MEMBER_VISIT",
|
||
"DWS_MEMBER_CONSUMPTION",
|
||
# 库存维度
|
||
"DWS_GOODS_STOCK_DAILY",
|
||
"DWS_GOODS_STOCK_WEEKLY",
|
||
"DWS_GOODS_STOCK_MONTHLY",
|
||
]
|
||
|
||
|
||
# ── 月窗口拆分 ────────────────────────────────────────────────
|
||
|
||
def split_by_month(start: date, end: date) -> list[tuple[date, date]]:
|
||
"""将 [start, end) 按自然月拆分为多个 (window_start, window_end) 窗口。
|
||
|
||
每个窗口：window_start = 当月1日, window_end = 次月1日（或 end，取较小值）。
|
||
"""
|
||
windows: list[tuple[date, date]] = []
|
||
cursor = start.replace(day=1) # 对齐到月初
|
||
if cursor < start:
|
||
cursor = start
|
||
while cursor < end:
|
||
# 次月1日
|
||
if cursor.month == 12:
|
||
next_month_first = date(cursor.year + 1, 1, 1)
|
||
else:
|
||
next_month_first = date(cursor.year, cursor.month + 1, 1)
|
||
window_end = min(next_month_first, end)
|
||
windows.append((cursor, window_end))
|
||
cursor = next_month_first
|
||
return windows
|
||
|
||
|
||
# ── 行数统计 ──────────────────────────────────────────────────
|
||
|
||
def _get_row_count_sql(task_code: str) -> str | None:
|
||
"""根据任务代码返回对应 DWS 表的行数查询 SQL。
|
||
|
||
返回 None 表示该任务无法直接统计行数（如 MV 刷新）。
|
||
"""
|
||
# 任务代码 → DWS 表名映射
|
||
table_map: dict[str, str] = {
|
||
"DWS_FINANCE_DAILY": "dws.dws_finance_daily_summary",
|
||
"DWS_FINANCE_RECHARGE": "dws.dws_finance_recharge_summary",
|
||
"DWS_FINANCE_INCOME_STRUCTURE": "dws.dws_finance_income_structure",
|
||
"DWS_FINANCE_DISCOUNT_DETAIL": "dws.dws_finance_discount_detail",
|
||
"DWS_ASSISTANT_DAILY": "dws.dws_assistant_daily_detail",
|
||
"DWS_ASSISTANT_ORDER_CONTRIBUTION": "dws.dws_assistant_order_contribution",
|
||
"DWS_ASSISTANT_CUSTOMER": "dws.dws_assistant_customer_daily",
|
||
"DWS_ASSISTANT_SALARY": "dws.dws_assistant_salary",
|
||
"DWS_ASSISTANT_FINANCE": "dws.dws_assistant_finance_analysis",
|
||
"DWS_ASSISTANT_MONTHLY": "dws.dws_assistant_monthly_summary",
|
||
"DWS_MEMBER_VISIT": "dws.dws_member_visit_daily",
|
||
"DWS_MEMBER_CONSUMPTION": "dws.dws_member_consumption_daily",
|
||
"DWS_GOODS_STOCK_DAILY": "dws.dws_goods_stock_daily",
|
||
"DWS_GOODS_STOCK_WEEKLY": "dws.dws_goods_stock_weekly",
|
||
"DWS_GOODS_STOCK_MONTHLY": "dws.dws_goods_stock_monthly",
|
||
}
|
||
table = table_map.get(task_code)
|
||
if not table:
|
||
return None
|
||
return f"SELECT COUNT(*) AS cnt FROM {table}"
|
||
|
||
|
||
def query_row_count(task_code: str, dsn: str) -> int | None:
|
||
"""查询指定任务对应 DWS 表的行数。需要 psycopg2。"""
|
||
sql = _get_row_count_sql(task_code)
|
||
if not sql:
|
||
return None
|
||
try:
|
||
import psycopg2
|
||
with psycopg2.connect(dsn) as conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute(sql)
|
||
row = cur.fetchone()
|
||
return row[0] if row else 0
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
# ── 单任务单窗口执行 ──────────────────────────────────────────
|
||
|
||
def run_task_window(
|
||
task_code: str,
|
||
window_start: date,
|
||
window_end: date,
|
||
dry_run: bool = False,
|
||
timeout_sec: int = 600,
|
||
) -> dict:
|
||
"""通过 CLI 子进程执行单个任务的单个月窗口。
|
||
|
||
返回 dict: task, window, success, elapsed, output
|
||
"""
|
||
cmd = [
|
||
"uv", "run", "--package", "etl-feiqiu",
|
||
"python", "-m", "cli.main",
|
||
"--layers", "DWS",
|
||
"--tasks", task_code,
|
||
"--window-start", f"{window_start} 00:00:00",
|
||
"--window-end", f"{window_end} 00:00:00",
|
||
"--force-full",
|
||
]
|
||
if dry_run:
|
||
cmd.append("--dry-run")
|
||
|
||
start_ts = time.time()
|
||
try:
|
||
proc = subprocess.run(
|
||
cmd,
|
||
cwd=str(ETL_CWD),
|
||
capture_output=True,
|
||
text=True,
|
||
encoding="utf-8",
|
||
errors="replace",
|
||
timeout=timeout_sec,
|
||
)
|
||
elapsed = time.time() - start_ts
|
||
output = proc.stdout + proc.stderr
|
||
has_error = any(
|
||
kw in line
|
||
for line in output.splitlines()
|
||
for kw in ("ERROR", "CRITICAL", "Traceback")
|
||
)
|
||
success = proc.returncode == 0 and not has_error
|
||
except subprocess.TimeoutExpired:
|
||
elapsed = time.time() - start_ts
|
||
output = f"TIMEOUT after {timeout_sec}s"
|
||
success = False
|
||
|
||
return {
|
||
"task": task_code,
|
||
"window": f"{window_start} ~ {window_end}",
|
||
"success": success,
|
||
"elapsed": elapsed,
|
||
"output": output,
|
||
}
|
||
|
||
|
||
# ── 主流程 ────────────────────────────────────────────────────
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
description="DWS 历史数据重算（按营业日口径）",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
)
|
||
parser.add_argument(
|
||
"--start-date", required=True,
|
||
help="重算起始日期 (YYYY-MM-DD)",
|
||
)
|
||
parser.add_argument(
|
||
"--end-date", required=True,
|
||
help="重算结束日期 (YYYY-MM-DD)",
|
||
)
|
||
parser.add_argument(
|
||
"--dry-run", action="store_true",
|
||
help="预览模式，不实际执行写入",
|
||
)
|
||
parser.add_argument(
|
||
"--fail-fast", action="store_true",
|
||
help="遇错即停（默认继续下一窗口）",
|
||
)
|
||
parser.add_argument(
|
||
"--tasks",
|
||
help="指定要重算的任务列表（逗号分隔），不指定则全部",
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
def main() -> int:
|
||
args = parse_args()
|
||
|
||
# 解析日期
|
||
try:
|
||
start_date = date.fromisoformat(args.start_date)
|
||
end_date = date.fromisoformat(args.end_date)
|
||
except ValueError as e:
|
||
print(f"日期格式错误: {e}", file=sys.stderr)
|
||
return 1
|
||
|
||
if start_date >= end_date:
|
||
print("错误: --start-date 必须早于 --end-date", file=sys.stderr)
|
||
return 1
|
||
|
||
# 确定任务列表
|
||
if args.tasks:
|
||
task_codes = [t.strip().upper() for t in args.tasks.split(",") if t.strip()]
|
||
# 校验任务代码合法性
|
||
invalid = [t for t in task_codes if t not in ALL_DWS_TASKS]
|
||
if invalid:
|
||
print(f"未知任务代码: {', '.join(invalid)}", file=sys.stderr)
|
||
print(f"可用任务: {', '.join(ALL_DWS_TASKS)}", file=sys.stderr)
|
||
return 1
|
||
else:
|
||
task_codes = list(ALL_DWS_TASKS)
|
||
|
||
# 拆分月窗口
|
||
windows = split_by_month(start_date, end_date)
|
||
|
||
# 日志目录
|
||
log_root = os.environ.get("SYSTEM_LOG_ROOT")
|
||
if not log_root:
|
||
raise RuntimeError(
|
||
"环境变量 SYSTEM_LOG_ROOT 未定义。"
|
||
"请在根 .env 中配置，参考 .env.template"
|
||
)
|
||
log_dir = Path(log_root)
|
||
log_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 数据库 DSN（用于行数统计）
|
||
pg_dsn = os.environ.get("PG_DSN", "")
|
||
|
||
now = datetime.now(TZ)
|
||
mode_label = "[DRY-RUN] " if args.dry_run else ""
|
||
fail_fast_label = " [FAIL-FAST]" if args.fail_fast else ""
|
||
|
||
print(f"{'='*70}")
|
||
print(f"{mode_label}DWS 历史数据重算 — 营业日口径{fail_fast_label}")
|
||
print(f"时间范围: {start_date} ~ {end_date}")
|
||
print(f"月窗口数: {len(windows)}")
|
||
print(f"任务数量: {len(task_codes)}")
|
||
print(f"任务列表: {', '.join(task_codes)}")
|
||
print(f"开始时间: {now.isoformat()}")
|
||
print(f"{'='*70}\n")
|
||
|
||
# 日志文件
|
||
log_file = log_dir / f"{now.strftime('%Y%m%d_%H%M%S')}_rebuild_dws_biz_date.log"
|
||
log_lines: list[str] = [
|
||
f"{mode_label}DWS 历史数据重算 — 营业日口径{fail_fast_label}\n",
|
||
f"时间范围: {start_date} ~ {end_date}\n",
|
||
f"月窗口数: {len(windows)}\n",
|
||
f"任务列表: {', '.join(task_codes)}\n",
|
||
f"开始时间: {now.isoformat()}\n",
|
||
f"{'='*70}\n\n",
|
||
]
|
||
|
||
# ── 执行前行数统计 ────────────────────────────────────────
|
||
row_counts_before: dict[str, int | None] = {}
|
||
if pg_dsn and not args.dry_run:
|
||
print("统计执行前行数...")
|
||
for tc in task_codes:
|
||
cnt = query_row_count(tc, pg_dsn)
|
||
row_counts_before[tc] = cnt
|
||
if cnt is not None:
|
||
log_lines.append(f"[BEFORE] {tc}: {cnt:,} 行\n")
|
||
log_lines.append("\n")
|
||
|
||
# ── 逐任务逐窗口执行 ──────────────────────────────────────
|
||
results: list[dict] = []
|
||
total_steps = len(task_codes) * len(windows)
|
||
step = 0
|
||
aborted = False
|
||
|
||
for task_code in task_codes:
|
||
if aborted:
|
||
break
|
||
for w_start, w_end in windows:
|
||
step += 1
|
||
label = f"[{step}/{total_steps}] {task_code} | {w_start} ~ {w_end}"
|
||
print(f"{label} ...", end=" ", flush=True)
|
||
|
||
result = run_task_window(
|
||
task_code, w_start, w_end,
|
||
dry_run=args.dry_run,
|
||
)
|
||
results.append(result)
|
||
|
||
icon = "OK" if result["success"] else "FAIL"
|
||
print(f"{icon} ({result['elapsed']:.0f}s)")
|
||
|
||
log_lines.append(f"[{icon}] {label} ({result['elapsed']:.0f}s)\n")
|
||
if not result["success"]:
|
||
# 输出最后 30 行错误信息
|
||
tail = result["output"].splitlines()[-30:]
|
||
for line in tail:
|
||
print(f" | {line}")
|
||
log_lines.append(f" | {line}\n")
|
||
log_lines.append("\n")
|
||
|
||
if args.fail_fast:
|
||
print("\n--fail-fast 已启用，中止后续执行")
|
||
log_lines.append("\n--fail-fast 中止\n")
|
||
aborted = True
|
||
break
|
||
|
||
# ── 执行后行数统计 ────────────────────────────────────────
|
||
row_counts_after: dict[str, int | None] = {}
|
||
if pg_dsn and not args.dry_run:
|
||
print("\n统计执行后行数...")
|
||
log_lines.append("\n")
|
||
for tc in task_codes:
|
||
cnt = query_row_count(tc, pg_dsn)
|
||
row_counts_after[tc] = cnt
|
||
before = row_counts_before.get(tc)
|
||
if cnt is not None:
|
||
diff = ""
|
||
if before is not None:
|
||
delta = cnt - before
|
||
sign = "+" if delta >= 0 else ""
|
||
diff = f" (差异: {sign}{delta:,})"
|
||
log_lines.append(f"[AFTER] {tc}: {cnt:,} 行{diff}\n")
|
||
|
||
# ── 汇总 ──────────────────────────────────────────────────
|
||
ok_count = sum(1 for r in results if r["success"])
|
||
fail_count = sum(1 for r in results if not r["success"])
|
||
total_elapsed = sum(r["elapsed"] for r in results)
|
||
|
||
summary = (
|
||
f"\n{'='*70}\n"
|
||
f"重算完成: {ok_count}/{len(results)} 成功, {fail_count} 失败"
|
||
f" (总耗时 {total_elapsed:.0f}s)\n"
|
||
)
|
||
print(summary)
|
||
log_lines.append(summary)
|
||
|
||
# 写入日志
|
||
with open(log_file, "w", encoding="utf-8") as f:
|
||
f.writelines(log_lines)
|
||
print(f"日志: {log_file}")
|
||
|
||
return 0 if fail_count == 0 else 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|