snake-python/scripts/migrate_gameplay_database.py

#!/usr/bin/env python3
"""Create a compact, normalized copy of a gameplay SQLite database.

The source is always opened read-only. The destination is written separately,
verified, and can optionally replace the source after a timestamped backup is
created. No in-place schema rewrite is performed.
"""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
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.backend.SqliteGameplayBackend import SqliteGameplayBackend
from server.database.game_quality import GameQualityInput, quality_meets_minimum, rate_game_quality

def open_source(path:Path) -> sqlite3.Connection:
  connection = sqlite3.connect(f"file:{path}?mode=ro", uri=True, timeout=60)
  connection.row_factory = sqlite3.Row
  connection.execute("PRAGMA query_only = ON")
  return connection

def object_columns(connection:sqlite3.Connection, name:str) -> set[str]:
  return {row[1] for row in connection.execute(f"PRAGMA table_info({name})")}

def column_or_null(columns:set[str], name:str) -> str:
  return name if name in columns else f"NULL AS {name}"

def create_destination(path:Path, busy_timeout_ms:int) -> sqlite3.Connection:
  if path.exists():
    raise FileExistsError(f"Destination already exists: {path}")
  SqliteGameplayBackend(
    str(path), busy_timeout_ms=busy_timeout_ms,
    initialize_indexes=False, journal_mode="DELETE",
  )
  connection = sqlite3.connect(str(path), timeout=max(1, busy_timeout_ms // 1000))
  connection.execute("PRAGMA foreign_keys = OFF")
  connection.execute("PRAGMA synchronous = OFF")
  connection.execute("PRAGMA locking_mode = EXCLUSIVE")
  connection.execute("PRAGMA temp_store = MEMORY")
  connection.execute("PRAGMA cache_size = -262144")
  connection.execute(f"PRAGMA busy_timeout = {busy_timeout_ms}")
  return connection

def analyze_quality(source:sqlite3.Connection) -> dict[str, object]:
  games = {
    row["game_id"]: row
    for row in source.execute("""
      SELECT game_id, status, final_turn, winner_name
      FROM games
    """)
  }
  aggregates = {
    row["game_id"]: row
    for row in source.execute("""
      SELECT t.game_id,
        COUNT(*) AS turn_rows,
        MIN(t.turn) AS min_turn,
        MAX(t.turn) AS max_turn,
        SUM(CASE WHEN t.my_move IN ('up','down','left','right') THEN 1 ELSE 0 END) AS valid_moves,
        SUM(CASE WHEN t.my_thinking_json IS NOT NULL AND t.my_thinking_json NOT IN ('', '{}', 'null') THEN 1 ELSE 0 END) AS thinking_rows,
        COUNT(DISTINCT t.my_move) AS distinct_moves
      FROM turns AS t
      GROUP BY t.game_id
    """)
  }
  snake_counts = {
    row["game_id"]: int(row["snake_turn_rows"])
    for row in source.execute("""
      SELECT game_id, COUNT(*) AS snake_turn_rows
      FROM snake_turns GROUP BY game_id
    """)
  }

  quality = {}
  for game_id, game in games.items():
    aggregate = aggregates.get(game_id)
    quality[game_id] = rate_game_quality(GameQualityInput(
      status=game["status"],
      final_turn=int(game["final_turn"] or 0),
      turn_rows=int(aggregate["turn_rows"] if aggregate else 0),
      min_turn=int(aggregate["min_turn"]) if aggregate and aggregate["min_turn"] is not None else None,
      max_turn=int(aggregate["max_turn"]) if aggregate and aggregate["max_turn"] is not None else None,
      valid_moves=int(aggregate["valid_moves"] if aggregate else 0),
      thinking_rows=int(aggregate["thinking_rows"] if aggregate else 0),
      distinct_moves=int(aggregate["distinct_moves"] if aggregate else 0),
      snake_turn_rows=snake_counts.get(game_id, 0),
      winner_name=game["winner_name"],
    ))
  return quality

def copy_games(source:sqlite3.Connection, destination:sqlite3.Connection, batch_size:int, quality:dict[str, object], minimum_quality:str) -> tuple[int, int]:
  columns = object_columns(source, "games")
  selected = [
    "game_id", "started_at", "ended_at", "width", "height", "source", "map_name",
    "ruleset_name", "ruleset_version", "your_snake_id", "your_snake_name",
    column_or_null(columns, "your_snake_type"), column_or_null(columns, "your_snake_version"),
    column_or_null(columns, "game_type"), column_or_null(columns, "winner_name"),
    "winner_you", "final_turn", "status",
  ]
  cursor = source.execute(f"SELECT {', '.join(selected)} FROM games ORDER BY game_id")
  sql = """
    INSERT INTO games (
      game_id, started_at, ended_at, width, height, source, map_name,
      ruleset_name, ruleset_version, your_snake_id, your_snake_name,
      your_snake_type, your_snake_version, game_type, winner_name,
      winner_you, final_turn, status, has_replay, quality_status,
      quality_score, quality_tier, quality_reasons_json
    ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
  """
  count = 0
  retained = 0
  while rows := cursor.fetchmany(batch_size):
    values = []
    for row in rows:
      game_quality = quality[row[0]]
      keep_replay = quality_meets_minimum(game_quality.tier, minimum_quality)
      values.append((
        *tuple(row), 1 if keep_replay else 0,
        "retained" if keep_replay else "low_quality",
        game_quality.score, game_quality.tier,
        json.dumps(game_quality.reasons, separators=(",", ":")),
      ))
      retained += int(keep_replay)
    destination.executemany(sql, values)
    count += len(rows)
  return count, retained

def retained_game_ids(quality:dict[str, object], minimum_quality:str) -> set[str]:
  return {
    game_id for game_id, result in quality.items()
    if quality_meets_minimum(result.tier, minimum_quality)
  }

def copy_game_snakes(source:sqlite3.Connection, destination:sqlite3.Connection, batch_size:int, retained_ids:set[str]) -> int:
  has_game_snakes = source.execute("""
    SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'game_snakes'
  """).fetchone() is not None
  if has_game_snakes:
    cursor = source.execute("""
      SELECT game_id, snake_id, snake_name, is_you
      FROM game_snakes ORDER BY game_id, snake_id
    """)
  else:
    cursor = source.execute("""
      SELECT game_id, snake_id, MAX(snake_name), MAX(is_you)
      FROM snake_turns
      GROUP BY game_id, snake_id
      ORDER BY game_id, snake_id
    """)
  sql = """
    INSERT INTO game_snakes (game_id, snake_id, snake_name, is_you)
    VALUES (?, ?, ?, ?)
  """
  count = 0
  while rows := cursor.fetchmany(batch_size):
    retained_rows = [tuple(row) for row in rows if row[0] in retained_ids]
    destination.executemany(sql, retained_rows)
    count += len(retained_rows)
  return count

def decode_json(value:str|None, fallback):
  if not value:
    return fallback
  try:
    return json.loads(value)
  except (json.JSONDecodeError, TypeError):
    return fallback

def compact_board(row:sqlite3.Row) -> tuple[str, str, str]:
  board = decode_json(row["board_state_json"], {})
  food = board.get("food") if isinstance(board, dict) else None
  hazards = board.get("hazards") if isinstance(board, dict) else None
  if food is None:
    food = decode_json(row["food_json"], [])
  if hazards is None:
    hazards = decode_json(row["hazards_json"], [])
  compact = json.dumps({}, separators=(",", ":"))
  return (
    compact,
    json.dumps(food or [], separators=(",", ":")),
    json.dumps(hazards or [], separators=(",", ":")),
  )

def copy_turns(source:sqlite3.Connection, destination:sqlite3.Connection, batch_size:int, retained_ids:set[str]) -> int:
  cursor = source.execute("""
    SELECT t.id, t.game_id, t.turn, t.observed_at, t.my_move, t.my_thinking_json,
      t.board_state_json, t.food_json, t.hazards_json
    FROM turns AS t
    ORDER BY t.id
  """)
  sql = """
    INSERT INTO turns (
      id, game_id, turn, observed_at, my_move, my_thinking_json,
      board_state_json, snakes_json, you_json, food_json, hazards_json
    ) VALUES (?,?,?,?,?,?,?,'[]','{}',?,?)
  """
  count = 0
  while rows := cursor.fetchmany(batch_size):
    values = []
    for row in rows:
      if row["game_id"] not in retained_ids:
        continue
      compact, food, hazards = compact_board(row)
      values.append((
        row["id"], row["game_id"], row["turn"], row["observed_at"],
        row["my_move"], row["my_thinking_json"], compact, food, hazards,
      ))
    destination.executemany(sql, values)
    count += len(values)
    if count % max(batch_size, 100_000) == 0:
      print(f"turns: {count:,}", flush=True)
  return count

def copy_snake_turns(source:sqlite3.Connection, destination:sqlite3.Connection, batch_size:int, retained_ids:set[str]) -> int:
  columns = object_columns(source, "snake_turns")
  latency = column_or_null(columns, "latency")
  cursor = source.execute(f"""
    SELECT st.id, st.game_id, st.turn, st.snake_id, st.health, st.length,
      st.head_x, st.head_y, st.body_json, st.inferred_move, st.{latency}
    FROM snake_turns AS st
    ORDER BY st.id
  """)
  sql = """
    INSERT INTO snake_turns (
      id, game_id, turn, snake_id, snake_name, health, length,
      head_x, head_y, body_json, is_you, inferred_move, latency
    ) VALUES (?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, 0, ?, ?)
  """
  count = 0
  while rows := cursor.fetchmany(batch_size):
    destination.executemany(sql, [
      (
        row["id"], row["game_id"], row["turn"], row["snake_id"],
        row["health"], row["length"], row["head_x"], row["head_y"],
        row["body_json"], row["inferred_move"], row["latency"],
      )
      for row in rows if row["game_id"] in retained_ids
    ])
    count += sum(1 for row in rows if row["game_id"] in retained_ids)
    if count % max(batch_size, 100_000) == 0:
      print(f"snake_turns: {count:,}", flush=True)
  return count

def verify(source:sqlite3.Connection, destination:sqlite3.Connection, retained_ids:set[str]) -> dict[str, int]:
  result = {}
  retained_turns = sum(
    1 for row in source.execute("SELECT game_id FROM turns")
    if row["game_id"] in retained_ids
  )
  retained_snake_turns = sum(
    1 for row in source.execute("SELECT game_id FROM snake_turns")
    if row["game_id"] in retained_ids
  )
  expected = {
    "games": int(source.execute("SELECT COUNT(*) FROM games").fetchone()[0]),
    "turns": retained_turns,
    "snake_turns": retained_snake_turns,
  }
  for table, source_count in expected.items():
    destination_count = int(destination.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0])
    if source_count != destination_count:
      raise RuntimeError(f"{table} count mismatch: {source_count} != {destination_count}")
    result[table] = destination_count
  integrity = destination.execute("PRAGMA integrity_check").fetchone()[0]
  if integrity != "ok":
    raise RuntimeError(f"Destination integrity check failed: {integrity}")
  return result

def replace_with_backup(source_path:Path, destination_path:Path) -> Path:
  timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
  backup = source_path.with_name(f"{source_path.name}.backup-{timestamp}")
  source_path.rename(backup)
  try:
    destination_path.rename(source_path)
  except Exception:
    backup.rename(source_path)
    raise
  return backup

def main() -> None:
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("--source", required=True, type=Path)
  parser.add_argument("--destination", type=Path)
  parser.add_argument("--batch-size", type=int, default=10_000)
  parser.add_argument("--busy-timeout-ms", type=int, default=60_000)
  parser.add_argument("--minimum-quality", choices=("low", "medium", "high"), default="medium", help="Minimum quality tier whose replay rows are retained")
  parser.add_argument("--replace", action="store_true", help="Backup source and replace it after verification")
  args = parser.parse_args()

  source_path = args.source.expanduser().resolve()
  destination_path = (args.destination or source_path.with_name(f"{source_path.stem}.compact{source_path.suffix}")).expanduser().resolve()
  if source_path == destination_path:
    raise SystemExit("Source and destination must be different paths")

  started = perf_counter()
  source = open_source(source_path)
  destination = create_destination(destination_path, max(1000, args.busy_timeout_ms))
  try:
    quality = analyze_quality(source)
    retained_ids = retained_game_ids(quality, args.minimum_quality)
    games, retained_games = copy_games(source, destination, max(1, args.batch_size), quality, args.minimum_quality)
    game_snakes = copy_game_snakes(source, destination, max(1, args.batch_size), retained_ids)
    turns = copy_turns(source, destination, max(1, args.batch_size), retained_ids)
    snake_turns = copy_snake_turns(source, destination, max(1, args.batch_size), retained_ids)
    destination.commit()
    destination.executescript("""
      CREATE INDEX IF NOT EXISTS idx_turns_game_turn ON turns(game_id, turn);
      CREATE INDEX IF NOT EXISTS idx_games_status ON games(status);
      CREATE INDEX IF NOT EXISTS idx_snake_turns_game_turn ON snake_turns(game_id, turn);
    """)
    destination.execute("PRAGMA foreign_keys = ON")
    counts = verify(source, destination, retained_ids)
  except Exception:
    destination.close()
    source.close()
    for suffix in ("", "-wal", "-shm"):
      Path(f"{destination_path}{suffix}").unlink(missing_ok=True)
    raise
  destination.close()
  source.close()

  source_bytes = source_path.stat().st_size
  destination_bytes = destination_path.stat().st_size
  print(f"copied: game rates={games:,}, retained replays={retained_games:,}, game_snakes={game_snakes:,}, turns={turns:,}, snake_turns={snake_turns:,}")
  print(f"verified: {counts}; size {source_bytes / 2**30:.2f} GiB -> {destination_bytes / 2**30:.2f} GiB")
  print(f"elapsed: {perf_counter() - started:.1f}s")

  if args.replace:
    backup = replace_with_backup(source_path, destination_path)
    print(f"source replaced; backup kept at: {backup}")
  else:
    print(f"compact database created at: {destination_path}")

if __name__ == "__main__":
  main()