ai-exchange/scripts/run_database_readiness.py

from pathlib import Path
from datetime import datetime, timezone
import json
import os
import sys
from typing import Optional

from sqlalchemy import text

ROOT = Path(__file__).resolve().parents[1]
os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://user:pass@localhost:5432/ai_ict_test")
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

DEFAULT_REQUIRED_VENUES = ("OKX",)
DEFAULT_REQUIRED_INSTRUMENTS = ("BTC-USDT",)
DEFAULT_MIN_CANDLES = 50
DEFAULT_MIN_CANDLE_SPAN_MINUTES = 60
DEFAULT_MAX_CANDLE_AGE_MINUTES = 240


def run_database_readiness(env=None, db_probe=None, schema_probe=None, seed_probe=None, candle_probe=None) -> dict:
    from scripts.run_autonomous_cycle import probe_database, probe_database_schema

    env = os.environ if env is None else env
    db_probe = db_probe or probe_database
    schema_probe = schema_probe or probe_database_schema
    seed_probe = seed_probe or probe_seed_data
    candle_probe = candle_probe or probe_candle_count
    required_venues = _get_csv(env, "READINESS_REQUIRED_VENUES", DEFAULT_REQUIRED_VENUES)
    required_instruments = _get_csv(env, "READINESS_REQUIRED_INSTRUMENTS", DEFAULT_REQUIRED_INSTRUMENTS)
    candle_instrument_id = _get_int(env, "READINESS_CANDLE_INSTRUMENT_ID", 1)
    candle_timeframe = env.get("READINESS_CANDLE_TIMEFRAME") or "1m"
    min_candles = _get_int(env, "READINESS_MIN_CANDLES", DEFAULT_MIN_CANDLES)
    min_candle_span_minutes = _get_int(env, "READINESS_MIN_CANDLE_SPAN_MINUTES", DEFAULT_MIN_CANDLE_SPAN_MINUTES)
    max_candle_age_minutes = _get_int(env, "READINESS_MAX_CANDLE_AGE_MINUTES", DEFAULT_MAX_CANDLE_AGE_MINUTES)
    now = datetime.now(timezone.utc)

    database = db_probe()
    if not database.get("available"):
        next_actions = build_next_actions(["database_unavailable"])
        return {
            "passed": False,
            "database": database,
            "schema": None,
            "seed_data": None,
            "candles": None,
            "reasons": ["database_unavailable"],
            "next_actions": next_actions,
            "action_plan": build_action_plan(next_actions),
        }

    schema = schema_probe()
    seed_data = seed_probe(required_venues=required_venues, required_instruments=required_instruments)
    candles = candle_probe(
        instrument_id=candle_instrument_id,
        timeframe=candle_timeframe,
        min_candles=min_candles,
        min_span_minutes=min_candle_span_minutes,
        max_age_minutes=max_candle_age_minutes,
        now=now,
    )
    reasons = []
    if not schema.get("ready"):
        reasons.append("schema_not_ready")
    if not seed_data.get("ready"):
        reasons.extend(seed_data.get("reasons") or ["seed_data_not_ready"])
    if not candles.get("ready"):
        reasons.extend(candles.get("reasons") or ["candles_not_ready"])

    next_actions = build_next_actions(reasons)
    return {
        "passed": len(reasons) == 0,
        "database": database,
        "schema": schema,
        "seed_data": seed_data,
        "candles": candles,
        "reasons": reasons,
        "next_actions": next_actions,
        "action_plan": build_action_plan(next_actions),
    }


def probe_seed_data(required_venues: tuple[str, ...] = DEFAULT_REQUIRED_VENUES, required_instruments: tuple[str, ...] = DEFAULT_REQUIRED_INSTRUMENTS) -> dict:
    from src.db.engine import engine

    with engine.connect() as connection:
        venue_codes = {
            code
            for code in required_venues
            if connection.execute(text("select code from venues where code = :code"), {"code": code}).scalar() is not None
        }
        instrument_symbols = {
            symbol
            for symbol in required_instruments
            if connection.execute(text("select symbol from instruments where symbol = :symbol"), {"symbol": symbol}).scalar() is not None
        }
    return evaluate_seed_readiness(
        found_venues=venue_codes,
        found_instruments=instrument_symbols,
        required_venues=required_venues,
        required_instruments=required_instruments,
    )


def probe_candle_count(
    instrument_id: int = 1,
    timeframe: str = "1m",
    min_candles: int = 0,
    min_span_minutes: int = 0,
    max_age_minutes: int = 0,
    now: Optional[datetime] = None,
) -> dict:
    from src.db.engine import engine

    with engine.connect() as connection:
        row = connection.execute(
            text(
                """
                select count(*) as candle_count, min(ts_open) as first_ts_open, max(ts_open) as last_ts_open
                from candles
                where instrument_id = :instrument_id and timeframe = :timeframe
                """
            ),
                {"instrument_id": instrument_id, "timeframe": timeframe},
        ).mappings().one()
    return evaluate_candle_readiness(
        candle_count=int(row["candle_count"] or 0),
        min_candles=min_candles,
        instrument_id=instrument_id,
        timeframe=timeframe,
        first_ts_open=row["first_ts_open"],
        last_ts_open=row["last_ts_open"],
        min_span_minutes=min_span_minutes,
        max_age_minutes=max_age_minutes,
        now=now,
    )


def evaluate_seed_readiness(found_venues: set[str], found_instruments: set[str], required_venues: tuple[str, ...], required_instruments: tuple[str, ...]) -> dict:
    missing_venues = [code for code in required_venues if code not in found_venues]
    missing_instruments = [symbol for symbol in required_instruments if symbol not in found_instruments]
    reasons = []
    if missing_venues:
        reasons.append("missing_required_venues")
    if missing_instruments:
        reasons.append("missing_required_instruments")
    return {
        "ready": len(reasons) == 0,
        "required_venues": list(required_venues),
        "required_instruments": list(required_instruments),
        "missing_venues": missing_venues,
        "missing_instruments": missing_instruments,
        "reasons": reasons,
    }


def evaluate_candle_readiness(
    candle_count: int,
    min_candles: int = 0,
    instrument_id: int = 1,
    timeframe: str = "1m",
    first_ts_open=None,
    last_ts_open=None,
    min_span_minutes: int = 0,
    max_age_minutes: int = 0,
    now: Optional[datetime] = None,
) -> dict:
    reasons = []
    if min_candles and candle_count < min_candles:
        reasons.append("not_enough_candles")
    span_minutes = compute_span_minutes(first_ts_open, last_ts_open)
    age_minutes = compute_age_minutes(last_ts_open, now=now)
    if min_span_minutes and (span_minutes is None or span_minutes < min_span_minutes):
        reasons.append("not_enough_candle_span")
    if max_age_minutes and (age_minutes is None or age_minutes > max_age_minutes):
        reasons.append("candles_too_stale")
    return {
        "ready": len(reasons) == 0,
        "instrument_id": instrument_id,
        "timeframe": timeframe,
        "candle_count": candle_count,
        "min_candles": min_candles,
        "first_ts_open": format_timestamp(first_ts_open),
        "last_ts_open": format_timestamp(last_ts_open),
        "span_minutes": span_minutes,
        "min_span_minutes": min_span_minutes,
        "age_minutes": age_minutes,
        "max_age_minutes": max_age_minutes,
        "reasons": reasons,
    }


def build_next_actions(reasons: list[str]) -> list[str]:
    actions = []
    if "database_unavailable" in reasons:
        actions.append("wait_for_database_recovery")
    if "schema_not_ready" in reasons:
        actions.append("run_migrations_before_pipeline")
    if "missing_required_venues" in reasons or "missing_required_instruments" in reasons:
        actions.append("verify_seed_migrations_or_seed_data")
    if "not_enough_candles" in reasons or "not_enough_candle_span" in reasons or "candles_too_stale" in reasons:
        actions.append("run_okx_backfill_or_load_replay_sample")
    if not actions:
        actions.append("run_guarded_pipeline")
    return actions


def build_action_plan(next_actions: list[str]) -> list[dict]:
    from scripts.run_autonomous_cycle import build_cycle_action_plan

    return build_cycle_action_plan(next_actions)


def compute_span_minutes(first_ts_open, last_ts_open) -> Optional[float]:
    first_ts = parse_timestamp(first_ts_open)
    last_ts = parse_timestamp(last_ts_open)
    if first_ts is None or last_ts is None:
        return None
    return max((last_ts - first_ts).total_seconds() / 60, 0.0)


def compute_age_minutes(last_ts_open, now: Optional[datetime] = None) -> Optional[float]:
    last_ts = parse_timestamp(last_ts_open)
    if last_ts is None:
        return None
    current = parse_timestamp(now or datetime.now(timezone.utc))
    return max((current - last_ts).total_seconds() / 60, 0.0)


def parse_timestamp(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


def format_timestamp(value):
    parsed = parse_timestamp(value)
    return parsed.isoformat() if parsed else None


def main(env=None) -> int:
    env = os.environ if env is None else env
    result = run_database_readiness(env=env)
    if _get_bool(env, "READINESS_OUTPUT_JSON", False):
        print(json.dumps(result, ensure_ascii=False, sort_keys=True))
    else:
        print(f"Database readiness passed: {result['passed']}")
        print(f"Database available: {result['database'].get('available')}")
        if result.get("schema") is not None:
            print(f"Schema ready: {result['schema'].get('ready')}")
        if result.get("seed_data") is not None:
            print(f"Seed data ready: {result['seed_data'].get('ready')}")
        if result.get("candles") is not None:
            print(f"Candles checked: {result['candles'].get('candle_count')}")
        if result["reasons"]:
            print(f"Database readiness reasons: {', '.join(result['reasons'])}")
        print(f"Database readiness next actions: {', '.join(result['next_actions'])}")
        for action in result.get("action_plan", []):
            command = " ".join(action["command"])
            db_hint = "db" if action["requires_database"] else "no-db"
            print(f"Database readiness action plan [{db_hint}]: {action['action']} -> {command}")
    return 0 if result["passed"] else 2


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"}


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