snake-python/scripts/migrate_gameplay_to_compact.py

#!/usr/bin/env python3
"""Migrate gameplay data from the legacy turns/snake_turns layout to compact storage.

Source and destination are a SQLite path or a postgresql:// DSN. Without
--destination the migration runs in place (same database, new tables). Every
game is read back and compared after writing; legacy tables are only removed
with --drop-legacy, and only when every legacy replay exists in compact form.

Online PostgreSQL procedure:
  1. run the migration while the old bot version keeps recording,
  2. deploy the new version (it records into the compact tables),
  3. run again with --since <start of step 1> to pick up games recorded meanwhile,
  4. run with --drop-legacy.
"""
import argparse
import asyncio
from pathlib import Path
import sqlite3
import sys
from time import perf_counter

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from server.database.benchmark_states import is_postgresql_source
from server.database.backend.PostgresqlGameplayBackend import PostgresqlGameplayBackend
from server.database.backend.SqliteGameplayBackend import SqliteGameplayBackend
from server.database.legacy.migrate import migrate_legacy
from server.database.legacy.reader_postgresql import LegacyPostgresqlReader
from server.database.legacy.reader_sqlite import LegacySqliteReader

_MISSING_REPLAYS = """
  SELECT COUNT(*) FROM games AS g
  WHERE g.has_replay
    AND EXISTS (SELECT 1 FROM turns AS t WHERE t.game_id = g.game_id)
    AND NOT EXISTS (SELECT 1 FROM game_replays AS r WHERE r.gid = g.gid)
    AND NOT EXISTS (SELECT 1 FROM game_turns AS c WHERE c.gid = g.gid)
"""

def _progress(done:int, total:int, stats) -> None:
  print(
    f"{done:,}/{total:,} games — replays={stats.replays:,} turns={stats.turns:,} "
    f"snake states={stats.states:,} failed={stats.failed:,}",
    flush=True,
  )

async def _drop_legacy(destination:str) -> None:
  if is_postgresql_source(destination):
    backend = PostgresqlGameplayBackend(destination)
    pool = await backend._get_pool()
    try:
      async with pool.acquire() as conn:
        missing = int(await conn.fetchval(_MISSING_REPLAYS))
        if missing:
          raise SystemExit(f"{missing} games still lack a compact replay; legacy tables kept")
        await conn.execute("DROP TABLE snake_turns")
        await conn.execute("DROP TABLE turns")
    finally:
      await backend.close()
  else:
    connection = sqlite3.connect(destination)
    try:
      missing = int(connection.execute(_MISSING_REPLAYS.replace("g.has_replay", "g.has_replay = 1")).fetchone()[0])
      if missing:
        raise SystemExit(f"{missing} games still lack a compact replay; legacy tables kept")
      connection.execute("DROP TABLE snake_turns")
      connection.execute("DROP TABLE turns")
      connection.commit()
      connection.execute("VACUUM")
    finally:
      connection.close()
  print("legacy tables turns and snake_turns dropped")

async def migrate(args:argparse.Namespace) -> None:
  source = args.source
  destination = args.destination or source
  if args.drop_legacy:
    await _drop_legacy(destination)
    return

  if is_postgresql_source(destination):
    backend = PostgresqlGameplayBackend(destination)
    await backend.initialize()
  else:
    backend = SqliteGameplayBackend(destination)

  if is_postgresql_source(source):
    reader = await LegacyPostgresqlReader.connect(source)
    close_reader = reader.close
  else:
    reader = LegacySqliteReader(source)
    async def close_reader() -> None:
      reader.close()

  started = perf_counter()
  try:
    stats = await migrate_legacy(
      reader, backend, since=args.since, minimum_quality=args.minimum_quality,
      verify=not args.no_verify, progress=_progress,
    )
  finally:
    await close_reader()
    await backend.close()

  print(
    f"migrated {stats.games:,} games ({stats.replays:,} replays, {stats.turns:,} turns, "
    f"{stats.states:,} snake states) in {perf_counter() - started:.1f}s"
  )
  if args.minimum_quality:
    print(f"replays dropped below '{args.minimum_quality}' quality: {stats.dropped_replays:,}")
  if stats.failed:
    print(f"FAILED games ({stats.failed}); their legacy rows are untouched:", file=sys.stderr)
    for failure in stats.failures[:50]:
      print(f"  {failure}", file=sys.stderr)
    raise SystemExit(1)

def main() -> None:
  parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
  parser.add_argument("--source", required=True, help="SQLite path or postgresql:// DSN with legacy tables")
  parser.add_argument("--destination", help="SQLite path or postgresql:// DSN (default: in place)")
  parser.add_argument("--since", help="Only games started at or after this ISO timestamp")
  parser.add_argument("--minimum-quality", choices=("low", "medium", "high"), help="Re-rate replays and drop those below this tier")
  parser.add_argument("--no-verify", action="store_true", help="Skip reading each game back after writing")
  parser.add_argument("--drop-legacy", action="store_true", help="Drop turns/snake_turns after checking every replay was migrated")
  asyncio.run(migrate(parser.parse_args()))

if __name__ == "__main__":
  main()