ai-exchange/scripts/run_okx_backfill.py

from pathlib import Path
from datetime import datetime, timezone
import json
import os
import sys
from http.client import IncompleteRead
from typing import Optional
from urllib.error import HTTPError, URLError

from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy import text

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from src.repositories.candle_repository import CandleRepository
from src.services.market_data.okx_client import OKXClient
from src.services.market_data.normalizer import normalize_okx_candles


INSTRUMENT_ID_MAP = {
    "BTC-USDT": 1,
    "ETH-USDT": 2,
}

TIMEFRAMES = ["1m", "5m", "15m"]


def build_backfill_config_from_env(env=None) -> dict:
    env = os.environ if env is None else env
    symbol = env.get("OKX_BACKFILL_SYMBOL") or "BTC-USDT"
    inst_id = env.get("OKX_BACKFILL_INST_ID") or f"{symbol}-SWAP"
    timeframes = _get_csv(env, "OKX_BACKFILL_TIMEFRAMES", tuple(TIMEFRAMES))
    return {
        "symbol": symbol,
        "inst_id": inst_id,
        "pages": _get_int(env, "OKX_BACKFILL_PAGES", 20),
        "limit": _get_int(env, "OKX_BACKFILL_LIMIT", 100),
        "timeframes": list(timeframes),
        "start_after": env.get("OKX_BACKFILL_AFTER") or None,
        "continue_from_earliest": _get_bool(env, "OKX_BACKFILL_CONTINUE_FROM_EARLIEST", False),
    }


def run_okx_backfill_once(
    client=None,
    repository=None,
    symbol: str = "BTC-USDT",
    inst_id: str = "BTC-USDT-SWAP",
    pages: int = 20,
    limit: int = 100,
    timeframes: list[str] = TIMEFRAMES,
    start_after: Optional[str] = None,
    continue_from_earliest: bool = False,
    initial_after_resolver=None,
) -> dict:
    if symbol not in INSTRUMENT_ID_MAP:
        raise ValueError(f"unsupported_okx_backfill_symbol:{symbol}")
    client = client or OKXClient()
    repository = repository or CandleRepository()
    initial_after_resolver = initial_after_resolver or resolve_initial_after_from_earliest
    grand_raw = 0
    grand_written = 0
    timeframe_results = []

    for timeframe in timeframes:
        total_raw = 0
        total_written = 0
        seen_ts = set()
        after = start_after
        if after is None and continue_from_earliest:
            after = initial_after_resolver(symbol=symbol, timeframe=timeframe)
        initial_after = after
        for _ in range(pages):
            raw = client.get_candles(inst_id=inst_id, bar=timeframe, limit=limit, after=after)
            if not raw:
                break
            normalized = normalize_okx_candles(symbol=symbol, venue="OKX", timeframe=timeframe, candles=raw)
            confirmed = []
            for candle in normalized:
                if not candle.confirm or candle.ts_open_ms in seen_ts:
                    continue
                seen_ts.add(candle.ts_open_ms)
                confirmed.append(candle)
            written = repository.upsert_many(confirmed, instrument_id=INSTRUMENT_ID_MAP[symbol])
            total_raw += len(raw)
            total_written += written
            next_after = str(min(int(candle.ts) for candle in raw))
            if next_after == after:
                break
            after = next_after
        grand_raw += total_raw
        grand_written += total_written
        timeframe_results.append({"timeframe": timeframe, "fetched": total_raw, "wrote": total_written, "initial_after": initial_after})

    return {
        "symbol": symbol,
        "inst_id": inst_id,
        "pages": pages,
        "limit": limit,
        "timeframes": timeframe_results,
        "total_fetched": grand_raw,
        "total_wrote": grand_written,
    }


def resolve_initial_after_from_earliest(symbol: str, timeframe: str) -> Optional[str]:
    instrument_id = INSTRUMENT_ID_MAP.get(symbol)
    if instrument_id is None:
        raise ValueError(f"unsupported_okx_backfill_symbol:{symbol}")

    from src.db.engine import engine

    with engine.connect() as connection:
        earliest = connection.execute(
            text(
                """
                select min(ts_open)
                from candles
                where instrument_id = :instrument_id and timeframe = :timeframe
                """
            ),
            {"instrument_id": instrument_id, "timeframe": timeframe},
        ).scalar()
    parsed = _parse_datetime(earliest)
    return str(int(parsed.timestamp() * 1000)) if parsed else None


def main(env=None) -> int:
    env = os.environ if env is None else env
    json_output = _get_bool(env, "OKX_BACKFILL_OUTPUT_JSON", False)
    try:
        summary = run_okx_backfill_once(**build_backfill_config_from_env(env=env))
    except (SQLAlchemyError, URLError, HTTPError, TimeoutError, RuntimeError, OSError, IncompleteRead, ValueError) as exc:
        summary = {
            "passed": False,
            "error_type": exc.__class__.__name__,
            "error": str(exc),
        }
        if json_output:
            print(json.dumps(summary, ensure_ascii=False, sort_keys=True))
        else:
            print(f"OKX backfill failed: {exc.__class__.__name__}")
        return 2
    summary["passed"] = True
    if json_output:
        print(json.dumps(summary, ensure_ascii=False, sort_keys=True))
        return 0
    for item in summary["timeframes"]:
        print(f"timeframe={item['timeframe']} fetched={item['fetched']} wrote={item['wrote']}")
    print(f"total fetched={summary['total_fetched']} total wrote={summary['total_wrote']}")
    return 0


def _get_csv(env, key: str, default: tuple[str, ...]) -> tuple[str, ...]:
    value = env.get(key)
    if value is None or value.strip() == "":
        return default
    return tuple(item.strip() for item in value.split(",") if item.strip())


def _get_int(env, key: str, default: int) -> int:
    value = env.get(key)
    if value is None or value == "":
        return default
    return int(value)


def _get_bool(env, key: str, default: bool) -> bool:
    value = env.get(key)
    if value is None or value == "":
        return default
    return str(value).strip().lower() in {"1", "true", "yes", "on"}


def _parse_datetime(value) -> Optional[datetime]:
    if value is None:
        return None
    if isinstance(value, datetime):
        if value.tzinfo is None:
            return value.replace(tzinfo=timezone.utc)
        return value.astimezone(timezone.utc)
    if isinstance(value, (int, float)):
        seconds = value / 1000 if value > 10_000_000_000 else value
        return datetime.fromtimestamp(seconds, tz=timezone.utc)
    if isinstance(value, str):
        normalized = value.strip().replace("Z", "+00:00")
        try:
            parsed = datetime.fromisoformat(normalized)
        except ValueError:
            return None
        if parsed.tzinfo is None:
            return parsed.replace(tzinfo=timezone.utc)
        return parsed.astimezone(timezone.utc)
    return None


if __name__ == "__main__":
    raise SystemExit(main())