"""扩展数据定时拉取引擎 — 从外部 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 # --------------------------------------------------------------------------- # 拉取执行 # --------------------------------------------------------------------------- def _apply_preset_flatten(config_id: str, rows: list[dict]) -> list[dict]: """对内置预设 (概念/行业) 应用结构转换, 与 fetch_preset 保持一致。 延迟导入避免与 ext_presets 形成循环依赖。 非预设 id 原样返回。 """ if config_id not in ("ext_gn_ths", "ext_hy_ths"): return rows from app.services.ext_presets import _flatten_concept_rows, _flatten_industry_rows flatten = _flatten_concept_rows if config_id == "ext_gn_ths" else _flatten_industry_rows return flatten(rows) 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") # 内置预设 (概念/行业): 应用结构转换, 让产出 schema 与分析页一致。 # 否则 raw 接口列 (concepts/industries 数组、name) 会直接覆盖正确的 part.parquet, # 导致分析页因找不到维度字段 (所属概念/所属同花顺行业) 而"数据消失"。 # 见 ext_presets._flatten_* —— 手动拉取 / 定时拉取都必须走同一套转换。 rows = _apply_preset_flatten(config.id, rows) # 字段映射 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 维护定时任务。 线程安全说明: refresh()/stop() 可能从主事件循环 (lifespan startup) 或同步路由的 worker 线程 (configure_pull 是 def 而非 async def, FastAPI 丢进线程池) 调用。worker 线程里没有 running loop, 直接 asyncio.create_task 会抛 "no running event loop"。因此对 task 的增删一律通过 call_soon_threadsafe 提交到主循环执行 —— 同一套代码两种调用场景都安全。 """ def __init__(self) -> None: self._tasks: dict[str, asyncio.Task] = {} self._running = False self._loop: asyncio.AbstractEventLoop | None = None def start(self, data_dir) -> None: """启动调度(在 lifespan startup 调用,主事件循环内)。""" self._running = True self._data_dir = data_dir try: self._loop = asyncio.get_running_loop() except RuntimeError: self._loop = None logger.info("PullScheduler started") def _submit(self, fn, *args) -> None: """把一个 callable 提交到主事件循环执行 (线程安全)。 startup 在主循环内调用时 fn 立即排队; worker 线程调用时跨线程排队。 两者都通过 call_soon_threadsafe, 保证 _tasks 字典的读写只在主循环里发生。 """ loop = self._loop if loop is None or loop.is_closed(): raise RuntimeError( "PullScheduler: 事件循环不可用 (start() 未在事件循环中调用?)" ) loop.call_soon_threadsafe(fn, *args) def stop(self) -> None: """停止所有任务 (从 shutdown 调用)。""" 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() new_configs: list[ExtConfig] = [] 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: new_configs.append(config) # 需要移除的 id (快照当前 task 字典的键, 避免遍历时改字典) remove_ids = [cid for cid in list(self._tasks) if cid not in active_ids] # 所有对 _tasks 的修改都提交到主循环里执行, 保证线程安全 def _apply() -> None: for config in new_configs: if config.id not in self._tasks: # 二次校验, 防重复 self._tasks[config.id] = self._loop.create_task( self._run_loop(config) ) logger.info( "PullScheduler: scheduled %s (every %d min)", config.id, config.pull.schedule_minutes, ) for cid in remove_ids: task = self._tasks.pop(cid, None) if task is not None: task.cancel() logger.info("PullScheduler: removed %s", cid) self._submit(_apply) async def _run_loop(self, config: ExtConfig) -> None: """单个配置的定时拉取循环。 策略: 启用后立即执行一次, 之后按 interval 循环。 每次循环重读最新配置 (fresh), interval 取自 fresh.pull.schedule_minutes, 这样用户中途修改间隔也能立即生效 (无需重启)。 """ try: while self._running: # 每轮重读最新配置 — 用户可能修改了 url / interval / enabled store = ExtConfigStore(self._data_dir) fresh = store.get(config.id) if not fresh or not fresh.pull or not fresh.pull.enabled: break pull = fresh.pull # 先执行一次 (启用即拉取, 让用户立刻看到生效) try: 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: fresh2 = store.get(config.id) if fresh2 and fresh2.pull: fresh2.pull.last_run = datetime.now(timezone.utc).isoformat() fresh2.pull.last_status = "error" fresh2.pull.last_message = str(e)[:200] store.upsert(fresh2) logger.warning("PullScheduler: %s error: %s", config.id, e) # 间隔取自最新配置 (每次重新读取, 修复改间隔不生效) interval = max(pull.schedule_minutes * 60, 60) # 至少 60s # 预告下次运行时间, 供前端展示 next_dt = datetime.now(timezone.utc).timestamp() + interval latest = store.get(config.id) if latest and latest.pull: latest.pull.next_run = datetime.fromtimestamp( next_dt, tz=timezone.utc ).isoformat() store.upsert(latest) await asyncio.sleep(interval) if not self._running: break except asyncio.CancelledError: pass async def _run_loop(self, config: ExtConfig) -> None: """单个配置的定时拉取循环。 策略: 启用后立即执行一次, 之后按 interval 循环。 每次循环重读最新配置 (fresh), interval 取自 fresh.pull.schedule_minutes, 这样用户中途修改间隔也能立即生效 (无需重启)。 """ try: while self._running: # 每轮重读最新配置 — 用户可能修改了 url / interval / enabled store = ExtConfigStore(self._data_dir) fresh = store.get(config.id) if not fresh or not fresh.pull or not fresh.pull.enabled: break pull = fresh.pull # 先执行一次 (启用即拉取, 让用户立刻看到生效) try: 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: fresh2 = store.get(config.id) if fresh2 and fresh2.pull: fresh2.pull.last_run = datetime.now(timezone.utc).isoformat() fresh2.pull.last_status = "error" fresh2.pull.last_message = str(e)[:200] store.upsert(fresh2) logger.warning("PullScheduler: %s error: %s", config.id, e) # 间隔取自最新配置 (每次重新读取, 修复改间隔不生效) interval = max(pull.schedule_minutes * 60, 60) # 至少 60s # 预告下次运行时间, 供前端展示 next_dt = datetime.now(timezone.utc).timestamp() + interval latest = store.get(config.id) if latest and latest.pull: latest.pull.next_run = datetime.fromtimestamp( next_dt, tz=timezone.utc ).isoformat() store.upsert(latest) await asyncio.sleep(interval) if not self._running: break except asyncio.CancelledError: pass # 全局单例 pull_scheduler = PullScheduler()