market_intel/apps/collector/main.py

"""采集入口：OKX 指数 + ATM Call/Put 周期采样落库。"""

from __future__ import annotations

import logging
import signal
import sys
import time
from typing import Any

from apps.collector.okx_rest import OkxRestClient
from apps.collector.selectors import rows_to_contracts, select_atm_pair
from packages.config import get_settings
from packages.db import Repository
from packages.db.repository import OptionQuoteRow
from packages.domain import option_leverage
from packages.notify import wecom

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s [collector] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("collector")

_STOP = False


def _handle_signal(signum: int, _frame: Any) -> None:
    global _STOP
    log.info("signal %s received, stopping…", signum)
    _STOP = True


def _now_ms() -> int:
    return int(time.time() * 1000)


def sample_once(
    client: OkxRestClient,
    repo: Repository,
    *,
    contracts_cache: list[dict[str, Any]],
    settings: Any,
) -> dict[str, Any]:
    ts_ms = _now_ms()
    index_px = client.fetch_index_ticker(settings.index_inst_id)
    if index_px is None or index_px <= 0:
        raise RuntimeError(f"index unavailable: {settings.index_inst_id}")

    repo.insert_index_tick(
        ts_ms=ts_ms,
        exchange="okx",
        underlying=settings.underlying,
        index_px=float(index_px),
    )

    pair = select_atm_pair(
        contracts_cache,
        index_px=float(index_px),
        min_hours=float(settings.min_option_hours),
    )
    if pair is None:
        raise RuntimeError("no eligible ATM option pair")

    meta: dict[str, Any] = {
        "index_px": index_px,
        "expiry_ymd": pair.expiry_ymd,
        "strike": pair.strike,
        "call_inst_id": pair.call_inst_id,
        "put_inst_id": pair.put_inst_id,
    }

    for side, inst_id in (("C", pair.call_inst_id), ("P", pair.put_inst_id)):
        ask, bid, ask_sz, bid_sz, book_ts = client.fetch_books(inst_id)
        lev = option_leverage(float(index_px), ask)
        repo.insert_option_quote(
            OptionQuoteRow(
                ts_ms=book_ts or ts_ms,
                exchange="okx",
                underlying=settings.underlying,
                inst_id=inst_id,
                expiry_ymd=pair.expiry_ymd,
                strike=pair.strike,
                side=side,
                index_px=float(index_px),
                ask=ask,
                bid=bid,
                ask_sz=ask_sz,
                bid_sz=bid_sz,
                leverage=lev,
            )
        )
        meta[f"{side}_ask"] = ask
        meta[f"{side}_leverage"] = lev

    return meta


def run() -> int:
    settings = get_settings()
    log.info(
        "start underlying=%s family=%s interval=%ss db=%s",
        settings.underlying,
        settings.option_inst_family,
        settings.sample_interval_sec,
        settings.db_path,
    )

    repo = Repository(settings.db_path)
    client = OkxRestClient(
        base_url=settings.okx_base_url,
        proxy=settings.okx_proxy or None,
    )

    contracts: list[dict[str, Any]] = []
    last_instruments_at = 0.0

    try:
        while not _STOP:
            t0 = time.monotonic()
            try:
                now = time.monotonic()
                if (
                    not contracts
                    or now - last_instruments_at >= float(settings.instruments_refresh_sec)
                ):
                    raw = client.fetch_option_instruments(settings.option_inst_family)
                    contracts = rows_to_contracts(raw)
                    last_instruments_at = now
                    log.info("instruments refreshed: %d contracts", len(contracts))

                meta = sample_once(client, repo, contracts_cache=contracts, settings=settings)
                repo.upsert_heartbeat(ok=True, meta=meta)
                wecom.notify_collector_recovered()
                log.info(
                    "sampled index=%.2f expiry=%s strike=%.0f C_lev=%s P_lev=%s",
                    meta["index_px"],
                    meta["expiry_ymd"],
                    meta["strike"],
                    f"{meta.get('C_leverage'):.1f}" if meta.get("C_leverage") else "-",
                    f"{meta.get('P_leverage'):.1f}" if meta.get("P_leverage") else "-",
                )
            except Exception as e:  # noqa: BLE001 — 单次失败记日志并跳过
                log.exception("sample failed: %s", e)
                repo.upsert_heartbeat(ok=False, error=str(e))
                hb = repo.get_heartbeat()
                wecom.notify_collector_fault(
                    error=str(e),
                    consecutive_failures=int(hb.get("consecutive_failures") or 0),
                )

            elapsed = time.monotonic() - t0
            sleep_for = max(1.0, float(settings.sample_interval_sec) - elapsed)
            # 可中断 sleep
            end = time.monotonic() + sleep_for
            while not _STOP and time.monotonic() < end:
                time.sleep(min(0.5, end - time.monotonic()))
    finally:
        client.close()
        repo.close()
        log.info("stopped")
    return 0


def main() -> None:
    signal.signal(signal.SIGINT, _handle_signal)
    signal.signal(signal.SIGTERM, _handle_signal)
    sys.exit(run())


if __name__ == "__main__":
    main()