trex-core/trex_core/runtime/reset.py

"""Asset-centric DB (re)initialization tool (dcy-020 / REQ-064, BAT-030).

Replaces the data-preserving pool→permanent cutover: instead of migrating
legacy pool rows, this **wipes the slot-core tables and (re)applies the
migrations** so the database lands on the clean asset-centric schema. Use it
to stand up a new instance's database or to reset a dev / superseded
pool-model database to the asset-centric substrate.

Why wipe instead of migrate (the decision that retired the cutover):
migration 009's partial `UNIQUE(instance_key, portfolio_id) WHERE portfolio_id
IS NOT NULL` cannot be applied to legacy pool data, because in the pool model
the same asset legitimately occupied more than one slot (a `closed` row from a
prior cycle plus a re-entered `holding` row) — so `portfolio_id` repeats and
the unique index collides. On an **empty** table the index applies cleanly and
`(instance_key, portfolio_id)` is unique by construction thereafter.

After reset, the daemon's universe reconcile (REQ-058) builds the permanent
asset slots (`inactive`); broker holdings surface through
`reconcile_core_holdings` (REQ-062) for operator adoption.

Destructive and DB-wide (every instance's rows). OPERATOR-RUN, dry-run by
default; pass `--yes` to apply. Take a backup first.
"""
from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Optional

_log = logging.getLogger(__name__)

# slot-core tables owned by trex-core migrations (drop order irrelevant — no FKs).
SLOT_CORE_TABLES: tuple[str, ...] = (
    "daily_slot",
    "daily_symbol_entries",
    "slot_transition_log",
    "intent_log",
    "ledger_outbox",
    "snapshot_pause",
    "killswitch_exclusion",
)


@dataclass
class ResetReport:
    db: str
    wiped: dict[str, int] = field(default_factory=dict)   # table -> rows removed
    migrations_applied: list[str] = field(default_factory=list)
    schema_ok: bool = False
    dry_run: bool = True

    def summary(self) -> str:
        total = sum(self.wiped.values())
        verb = "would wipe" if self.dry_run else "wiped"
        return (
            f"reset[{self.db}] {verb} {total} rows across {len(self.wiped)} tables "
            f"{self.wiped} | migrations={len(self.migrations_applied)} "
            f"schema_ok={self.schema_ok} dry_run={self.dry_run}"
        )


def _existing_tables(cur) -> set[str]:
    cur.execute(
        "SELECT table_name FROM information_schema.tables "
        "WHERE table_schema = 'public' AND table_name = ANY(%s)",
        (list(SLOT_CORE_TABLES),),
    )
    return {r[0] for r in cur.fetchall()}


def run_reset(
    dsn: str,
    *,
    dry_run: bool = True,
    psycopg_module: Any = None,
) -> ResetReport:
    """Wipe the slot-core tables and reapply migrations to the asset-centric schema.

    dry_run (default): report the row counts that WOULD be wiped, change nothing.
    apply (dry_run=False): TRUNCATE existing slot-core tables, then
    `apply_migrations(dsn)` so the schema is the clean asset-centric form, then
    verify `daily_slot` columns == `EXPECTED_SLOT_FIELDS`.
    """
    if psycopg_module is None:
        import psycopg  # type: ignore[import-not-found]
        psycopg_module = psycopg

    with psycopg_module.connect(dsn, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT current_database()")
            report = ResetReport(db=cur.fetchone()[0], dry_run=dry_run)
            existing = _existing_tables(cur)
            for t in SLOT_CORE_TABLES:
                if t in existing:
                    cur.execute(f"SELECT count(*) FROM {t}")
                    report.wiped[t] = int(cur.fetchone()[0])
            if dry_run:
                return report
            if existing:
                cur.execute(
                    f"TRUNCATE {', '.join(existing)} RESTART IDENTITY CASCADE"
                )

    # Reapply migrations on the now-empty DB: 009's unique index applies cleanly,
    # 010 sets the inactive CHECK, 011 drops the pool-model columns.
    from trex_core.adapters.persistence.migrate import apply_migrations

    report.migrations_applied = apply_migrations(dsn, psycopg_module=psycopg_module)

    from trex_core.domain.slot import EXPECTED_SLOT_FIELDS

    with psycopg_module.connect(dsn) as conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT column_name FROM information_schema.columns "
                "WHERE table_name = 'daily_slot'"
            )
            cols = {r[0] for r in cur.fetchall()}
    report.schema_ok = cols == set(EXPECTED_SLOT_FIELDS)
    return report


def main(argv: Optional[list[str]] = None) -> int:  # pragma: no cover - CLI shell
    import argparse

    from trex_core.runtime.config import DaemonConfig, resolve_config_toml

    logging.basicConfig(level=logging.INFO)
    parser = argparse.ArgumentParser(
        prog="trex-core-reset",
        description="dcy-020/REQ-064: wipe the slot-core tables and reinitialize "
                    "the DB to the asset-centric schema (dry-run by default). "
                    "Destructive and DB-wide — back up first.",
    )
    parser.add_argument("--config", required=True, help="instance TOML name/path (DSN source)")
    parser.add_argument("--config-path", default=None)
    parser.add_argument("--env-file", default=None)
    parser.add_argument("--yes", action="store_true",
                        help="actually wipe + reinitialize (otherwise dry-run)")
    args = parser.parse_args(argv)

    toml_path = resolve_config_toml(args.config, args.config_path)
    config = DaemonConfig.from_config(toml_path, env_file=args.env_file)

    report = run_reset(config.dsn, dry_run=not args.yes)
    _log.warning("%s", report.summary())
    if not args.yes:
        _log.info("dry-run only — re-run with --yes to wipe and reinitialize")
    elif not report.schema_ok:
        _log.error("schema verification FAILED after reset — inspect manually")
        return 1
    return 0


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())