@@ -0,0 +1,200 @@
|
||||
"""数据同步 — 将本地 data/ 打包推送到 serve 端。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import tarfile
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api/data/sync", tags=["data-sync"])
|
||||
|
||||
# 允许同步的子目录 — 覆盖 serve 功能面板所需的数据
|
||||
SYNCABLE_PARTS = [
|
||||
"kline_daily",
|
||||
"kline_daily_enriched",
|
||||
"kline_index_daily",
|
||||
"kline_index_enriched",
|
||||
"kline_etf_daily",
|
||||
"kline_etf_enriched",
|
||||
"kline_etf_minute",
|
||||
"kline_minute",
|
||||
"adj_factor",
|
||||
"adj_factor_etf",
|
||||
"instruments",
|
||||
"instruments_index",
|
||||
"instruments_etf",
|
||||
"financials",
|
||||
"pools",
|
||||
]
|
||||
|
||||
|
||||
def _part_size_info(data_dir: Path, part: str) -> dict:
|
||||
"""获取指定子目录的文件数和大小。"""
|
||||
part_dir = data_dir / part
|
||||
if not part_dir.exists():
|
||||
return {"file_count": 0, "size_bytes": 0, "last_modified": None}
|
||||
files = list(part_dir.rglob("*.parquet"))
|
||||
if not files:
|
||||
return {"file_count": 0, "size_bytes": 0, "last_modified": None}
|
||||
total = sum(f.stat().st_size for f in files)
|
||||
mtimes = [f.stat().st_mtime for f in files]
|
||||
last_ts = datetime.fromtimestamp(max(mtimes)).isoformat() if mtimes else None
|
||||
return {"file_count": len(files), "size_bytes": total, "last_modified": last_ts}
|
||||
|
||||
|
||||
def _pack_parts(data_dir: Path, parts: list[str]) -> bytes:
|
||||
"""将指定子目录打包为 tar.gz 字节流。"""
|
||||
buf = BytesIO()
|
||||
with tarfile.open(fileobj=buf, mode="w:gz") as tar:
|
||||
for part in parts:
|
||||
part_dir = data_dir / part
|
||||
if not part_dir.exists():
|
||||
continue
|
||||
for fpath in sorted(part_dir.rglob("*.parquet")):
|
||||
rel = fpath.relative_to(data_dir)
|
||||
tar.add(fpath, arcname=str(rel))
|
||||
buf.seek(0)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
class SyncConfig(BaseModel):
|
||||
serve_url: str = ""
|
||||
sync_key: str = ""
|
||||
enable_auto: bool = False
|
||||
interval_minutes: int = 0
|
||||
|
||||
|
||||
# ── 存储配置用的 JSON 文件 ──────────────────────────────────
|
||||
CONFIG_FILE = "sync_config.json"
|
||||
|
||||
|
||||
def _config_path(data_dir: Path) -> Path:
|
||||
return data_dir / "user_data" / CONFIG_FILE
|
||||
|
||||
|
||||
def _load_config(data_dir: Path) -> SyncConfig:
|
||||
cp = _config_path(data_dir)
|
||||
if cp.exists():
|
||||
try:
|
||||
return SyncConfig(**json.loads(cp.read_text()))
|
||||
except Exception:
|
||||
pass
|
||||
return SyncConfig()
|
||||
|
||||
|
||||
def _save_config(data_dir: Path, cfg: SyncConfig) -> None:
|
||||
cp = _config_path(data_dir)
|
||||
cp.parent.mkdir(parents=True, exist_ok=True)
|
||||
cp.write_text(cfg.model_dump_json(indent=2))
|
||||
|
||||
|
||||
# ── 路由 ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
def sync_status(request: Request):
|
||||
"""各数据目录的本地状态 + 同步配置。"""
|
||||
data_dir: Path = request.app.state.datastore.data_dir
|
||||
cfg = _load_config(data_dir)
|
||||
parts = {}
|
||||
for p in SYNCABLE_PARTS:
|
||||
parts[p] = _part_size_info(data_dir, p)
|
||||
return {
|
||||
"parts": parts,
|
||||
"config": cfg.model_dump(),
|
||||
"last_sync": cfg.model_dump().get("_last_sync_time"),
|
||||
}
|
||||
|
||||
|
||||
class ConfigReq(BaseModel):
|
||||
serve_url: str
|
||||
sync_key: str
|
||||
enable_auto: bool = False
|
||||
interval_minutes: int = 0
|
||||
|
||||
|
||||
@router.post("/config")
|
||||
def update_sync_config(request: Request, body: ConfigReq):
|
||||
"""更新同步配置。"""
|
||||
data_dir: Path = request.app.state.datastore.data_dir
|
||||
cfg = SyncConfig(
|
||||
serve_url=body.serve_url.rstrip("/"),
|
||||
sync_key=body.sync_key,
|
||||
enable_auto=body.enable_auto,
|
||||
interval_minutes=body.interval_minutes,
|
||||
)
|
||||
_save_config(data_dir, cfg)
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
class PushReq(BaseModel):
|
||||
parts: list[str]
|
||||
|
||||
|
||||
@router.post("/push")
|
||||
async def push_sync(request: Request, body: PushReq):
|
||||
"""打包指定 parts 推送到 serve 端。"""
|
||||
data_dir: Path = request.app.state.datastore.data_dir
|
||||
cfg = _load_config(data_dir)
|
||||
if not cfg.serve_url:
|
||||
raise HTTPException(status_code=400, detail="未配置 serve_url,请先在同步设置中配置")
|
||||
|
||||
requested = [p for p in body.parts if p in SYNCABLE_PARTS]
|
||||
if not requested:
|
||||
raise HTTPException(status_code=400, detail="未指定有效的同步目录")
|
||||
|
||||
# 检查有没有数据可同步
|
||||
available = [p for p in requested if (data_dir / p).exists()]
|
||||
if not available:
|
||||
raise HTTPException(status_code=400, detail="所选目录均无数据可同步")
|
||||
|
||||
# 打包
|
||||
tarball = _pack_parts(data_dir, available)
|
||||
size_bytes = len(tarball)
|
||||
logger.info("packed %d parts -> %d bytes", len(available), size_bytes)
|
||||
|
||||
# 推送
|
||||
upload_url = f"{cfg.serve_url}/api/data/sync/upload"
|
||||
headers = {"X-Sync-Key": cfg.sync_key} if cfg.sync_key else {}
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=600) as client:
|
||||
resp = await client.post(
|
||||
upload_url,
|
||||
files={"file": ("sync.tar.gz", tarball, "application/gzip")},
|
||||
data={"parts": ",".join(available)},
|
||||
headers=headers,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
detail = resp.text[:500]
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail=f"serve 返回异常 ({resp.status_code}): {detail}",
|
||||
)
|
||||
result = resp.json()
|
||||
except httpx.RequestError as exc:
|
||||
raise HTTPException(status_code=502, detail=f"连接 serve 失败: {exc}") from exc
|
||||
|
||||
# 记录上次同步时间
|
||||
cfg_dict = cfg.model_dump()
|
||||
cfg_dict["_last_sync_time"] = datetime.now().isoformat()
|
||||
cfg = SyncConfig(**cfg_dict)
|
||||
_save_config(data_dir, cfg)
|
||||
|
||||
logger.info("sync pushed: parts=%s files=%d", available, result.get("file_count", 0))
|
||||
return {
|
||||
"ok": True,
|
||||
"parts_sent": available,
|
||||
"size_bytes": size_bytes,
|
||||
"serve_response": result,
|
||||
}
|
||||
Reference in New Issue
Block a user