初始化工程
This commit is contained in:
@@ -0,0 +1,216 @@
|
||||
"""扩展数据定时拉取引擎 — 从外部 API 拉取数据写入 Parquet。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from datetime import date, datetime, timezone
|
||||
from functools import reduce
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.services.ext_data import (
|
||||
ExtConfig,
|
||||
ExtConfigStore,
|
||||
PullConfig,
|
||||
rows_to_parquet,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 响应解析
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _extract_rows(data: Any, path: str) -> list[dict]:
|
||||
"""按 dot-path 从 JSON 响应中提取行数组。
|
||||
|
||||
例: path="data.list" → response["data"]["list"]
|
||||
如果 path 为空，直接将 data 视为数组。
|
||||
"""
|
||||
if not path:
|
||||
if isinstance(data, list):
|
||||
return data
|
||||
raise ValueError("response_path 为空但响应不是数组")
|
||||
|
||||
keys = path.split(".")
|
||||
current = data
|
||||
for key in keys:
|
||||
if isinstance(current, dict):
|
||||
if key not in current:
|
||||
raise ValueError(f"响应中不存在路径 '{path}'，缺失键 '{key}'")
|
||||
current = current[key]
|
||||
elif isinstance(current, list):
|
||||
try:
|
||||
current = current[int(key)]
|
||||
except (ValueError, IndexError) as e:
|
||||
raise ValueError(f"响应路径 '{path}' 解析失败: {e}") from e
|
||||
else:
|
||||
raise ValueError(f"响应路径 '{path}' 中间值不是 dict/list: {type(current)}")
|
||||
|
||||
if not isinstance(current, list):
|
||||
raise ValueError(f"路径 '{path}' 指向的不是数组，而是 {type(current)}")
|
||||
|
||||
return current
|
||||
|
||||
|
||||
def _apply_field_map(rows: list[dict], field_map: dict[str, str]) -> list[dict]:
|
||||
"""将外部字段名映射为内部配置字段名。field_map: {外部名: 内部名}。"""
|
||||
if not field_map:
|
||||
return rows
|
||||
mapped = []
|
||||
for row in rows:
|
||||
new_row: dict = {}
|
||||
for k, v in row.items():
|
||||
mapped_key = field_map.get(k, k)
|
||||
new_row[mapped_key] = v
|
||||
mapped.append(new_row)
|
||||
return mapped
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 拉取执行
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def fetch_and_ingest(
|
||||
config: ExtConfig,
|
||||
data_dir,
|
||||
) -> tuple[int, str]:
|
||||
"""执行一次拉取: 请求外部 API → 解析响应 → 写入 Parquet。
|
||||
|
||||
Returns:
|
||||
(rows_written, date_str)
|
||||
"""
|
||||
pull = config.pull
|
||||
if not pull or not pull.url:
|
||||
raise ValueError("拉取未配置或 URL 为空")
|
||||
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
headers = pull.headers or {}
|
||||
kwargs: dict[str, Any] = {"headers": headers}
|
||||
|
||||
if pull.method.upper() == "POST" and pull.body:
|
||||
kwargs["content"] = pull.body
|
||||
if "content-type" not in {k.lower() for k in headers}:
|
||||
kwargs["headers"]["Content-Type"] = "application/json"
|
||||
|
||||
resp = await client.request(pull.method.upper(), pull.url, **kwargs)
|
||||
resp.raise_for_status()
|
||||
|
||||
# 解析 JSON
|
||||
try:
|
||||
data = resp.json()
|
||||
except Exception as e:
|
||||
raise ValueError(f"响应不是有效 JSON: {e}") from e
|
||||
|
||||
# 提取行
|
||||
rows = _extract_rows(data, pull.response_path)
|
||||
if not rows:
|
||||
raise ValueError("提取到的行数为 0")
|
||||
|
||||
# 字段映射
|
||||
rows = _apply_field_map(rows, pull.field_map)
|
||||
|
||||
# 校验 symbol 列
|
||||
if rows and "symbol" not in rows[0]:
|
||||
raise ValueError("数据行中缺少 symbol 字段，请配置 field_map 映射")
|
||||
|
||||
# 写入
|
||||
snap = date.today()
|
||||
n = rows_to_parquet(rows, config, data_dir, snapshot_date=snap)
|
||||
return n, snap.isoformat()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 调度器
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class PullScheduler:
|
||||
"""后台调度器：为每个启用了 pull 的 ExtConfig 维护定时任务。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._tasks: dict[str, asyncio.Task] = {}
|
||||
self._running = False
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def start(self, data_dir) -> None:
|
||||
"""启动调度（在 lifespan startup 调用）。"""
|
||||
self._running = True
|
||||
self._data_dir = data_dir
|
||||
logger.info("PullScheduler started")
|
||||
|
||||
def stop(self) -> None:
|
||||
"""停止所有任务。"""
|
||||
self._running = False
|
||||
for task in self._tasks.values():
|
||||
task.cancel()
|
||||
self._tasks.clear()
|
||||
logger.info("PullScheduler stopped")
|
||||
|
||||
def refresh(self, data_dir) -> None:
|
||||
"""重新加载配置，更新调度任务（增/删/改）。"""
|
||||
self._data_dir = data_dir
|
||||
store = ExtConfigStore(data_dir)
|
||||
configs = store.load_all()
|
||||
|
||||
active_ids: set[str] = set()
|
||||
|
||||
for config in configs:
|
||||
if not config.pull or not config.pull.enabled or not config.pull.url:
|
||||
continue
|
||||
active_ids.add(config.id)
|
||||
if config.id not in self._tasks:
|
||||
# 新增调度
|
||||
task = asyncio.create_task(self._run_loop(config))
|
||||
self._tasks[config.id] = task
|
||||
logger.info("PullScheduler: scheduled %s (every %d min)", config.id, config.pull.schedule_minutes)
|
||||
|
||||
# 移除不再活跃的
|
||||
for cid in list(self._tasks):
|
||||
if cid not in active_ids:
|
||||
self._tasks[cid].cancel()
|
||||
del self._tasks[cid]
|
||||
logger.info("PullScheduler: removed %s", cid)
|
||||
|
||||
async def _run_loop(self, config: ExtConfig) -> None:
|
||||
"""单个配置的定时拉取循环。"""
|
||||
try:
|
||||
while self._running:
|
||||
pull = config.pull
|
||||
if not pull:
|
||||
break
|
||||
interval = max(pull.schedule_minutes * 60, 60) # 至少 60s
|
||||
await asyncio.sleep(interval)
|
||||
if not self._running:
|
||||
break
|
||||
try:
|
||||
# 重新加载最新配置（用户可能中途修改）
|
||||
store = ExtConfigStore(self._data_dir)
|
||||
fresh = store.get(config.id)
|
||||
if not fresh or not fresh.pull or not fresh.pull.enabled:
|
||||
break
|
||||
n, d = await fetch_and_ingest(fresh, self._data_dir)
|
||||
fresh.pull.last_run = datetime.now(timezone.utc).isoformat()
|
||||
fresh.pull.last_status = "success"
|
||||
fresh.pull.last_message = f"{n} rows @ {d}"
|
||||
fresh.pull.last_rows = n
|
||||
store.upsert(fresh)
|
||||
logger.info("PullScheduler: %s success, %d rows", config.id, n)
|
||||
except Exception as e:
|
||||
store = ExtConfigStore(self._data_dir)
|
||||
fresh = store.get(config.id)
|
||||
if fresh and fresh.pull:
|
||||
fresh.pull.last_run = datetime.now(timezone.utc).isoformat()
|
||||
fresh.pull.last_status = "error"
|
||||
fresh.pull.last_message = str(e)[:200]
|
||||
store.upsert(fresh)
|
||||
logger.warning("PullScheduler: %s error: %s", config.id, e)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
|
||||
# 全局单例
|
||||
pull_scheduler = PullScheduler()
|
||||
Reference in New Issue
Block a user