snake-python/scripts/merge_gameplay_databases.py

#!/usr/bin/env python3
"""Safely merge a cleaned gameplay database with a temporary delta database.

The base and delta are opened read-only. A new destination is created, numeric
row IDs are regenerated, delta games are quality-rated, and all result rows are
kept even when their replay is excluded. The base can be replaced only after
verification and a timestamped backup.
"""

from __future__ import annotations

import argparse
from collections import Counter
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 scripts.migrate_gameplay_database import (
  analyze_quality, column_or_null, create_destination, object_columns,
  open_source, retained_game_ids,
)
from server.database.game_quality import quality_meets_minimum

GAME_COLUMNS = (
  "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",
)
GAME_DEFAULTS = {
  "winner_you": "0 AS winner_you", "final_turn": "0 AS final_turn",
  "status": "'running' AS status", "has_replay": "1 AS has_replay",
  "quality_status": "'retained' AS quality_status",
}
CORE_GAME_FIELDS = (
  "game_id", "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",
)

def select_expression(columns:set[str], name:str) -> str:
  if name in columns:
    return name
  return GAME_DEFAULTS.get(name, f"NULL AS {name}")

def read_games(connection:sqlite3.Connection) -> dict[str, sqlite3.Row]:
  columns = object_columns(connection, "games")
  selected = ", ".join(select_expression(columns, name) for name in GAME_COLUMNS)
  return {
    row["game_id"]: row
    for row in connection.execute(f"SELECT {selected} FROM games")
  }

def duplicate_ids(base_games:dict[str, sqlite3.Row], delta_games:dict[str, sqlite3.Row]) -> set[str]:
  duplicates = set(base_games) & set(delta_games)
  conflicts = []
  for game_id in duplicates:
    base_signature = tuple(base_games[game_id][name] for name in CORE_GAME_FIELDS)
    delta_signature = tuple(delta_games[game_id][name] for name in CORE_GAME_FIELDS)
    if base_signature != delta_signature:
      conflicts.append(game_id)
  if conflicts:
    examples = ", ".join(sorted(conflicts)[:5])
    raise RuntimeError(
      f"Conflicting duplicate game_id values ({len(conflicts)}): {examples}. "
      "Nothing was merged."
    )
  return duplicates

def insert_games(
  destination:sqlite3.Connection,
  rows:dict[str, sqlite3.Row],
  excluded:set[str],
  quality:dict[str, object]|None=None,
  minimum_quality:str="medium",
) -> tuple[int, set[str], Counter]:
  placeholders = ",".join("?" for _ in GAME_COLUMNS)
  sql = f"INSERT INTO games ({','.join(GAME_COLUMNS)}) VALUES ({placeholders})"
  values = []
  replay_ids:set[str] = set()
  tiers:Counter = Counter()

  for game_id, row in rows.items():
    if game_id in excluded:
      continue
    output = [row[name] for name in GAME_COLUMNS]
    if quality is not None:
      result = quality[game_id]
      keep_replay = quality_meets_minimum(result.tier, minimum_quality)
      replacements = {
        "has_replay": int(keep_replay),
        "quality_status": "retained" if keep_replay else "low_quality",
        "quality_score": result.score,
        "quality_tier": result.tier,
        "quality_reasons_json": json.dumps(result.reasons, separators=(",", ":")),
      }
      output = [replacements.get(name, row[name]) for name in GAME_COLUMNS]
      tiers[result.tier] += 1
      if keep_replay:
        replay_ids.add(game_id)
    elif bool(row["has_replay"]):
      replay_ids.add(game_id)
    values.append(tuple(output))

  destination.executemany(sql, values)
  return len(values), replay_ids, tiers

def copy_game_snakes(source:sqlite3.Connection, destination:sqlite3.Connection, allowed:set[str], batch_size:int) -> int:
  has_table = source.execute(
    "SELECT 1 FROM sqlite_master WHERE type='table' AND name='game_snakes'"
  ).fetchone()
  if has_table:
    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):
    values = [tuple(row) for row in rows if row[0] in allowed]
    destination.executemany(sql, values)
    count += len(values)
  return count

def copy_turns(source:sqlite3.Connection, destination:sqlite3.Connection, allowed:set[str], batch_size:int) -> int:
  columns = object_columns(source, "turns")
  names = (
    "game_id", "turn", "observed_at", "my_move", "my_thinking_json",
    "board_state_json", "snakes_json", "you_json", "food_json", "hazards_json",
  )
  defaults = {
    "my_thinking_json": "NULL AS my_thinking_json", "board_state_json": "'{}' AS board_state_json",
    "snakes_json": "'[]' AS snakes_json", "you_json": "'{}' AS you_json",
    "food_json": "'[]' AS food_json", "hazards_json": "'[]' AS hazards_json",
  }
  selected = ",".join(name if name in columns else defaults[name] for name in names)
  cursor = source.execute(f"SELECT {selected} FROM turns ORDER BY id")
  sql = f"INSERT INTO turns ({','.join(names)}) VALUES ({','.join('?' for _ in names)})"
  count = 0
  while rows := cursor.fetchmany(batch_size):
    values = [tuple(row) for row in rows if row[0] in allowed]
    destination.executemany(sql, values)
    count += len(values)
  return count

def copy_snake_turns(source:sqlite3.Connection, destination:sqlite3.Connection, allowed:set[str], batch_size:int) -> int:
  columns = object_columns(source, "snake_turns")
  names = (
    "game_id", "turn", "snake_id", "snake_name", "health", "length",
    "head_x", "head_y", "body_json", "is_you", "inferred_move", "latency",
  )
  selected = ",".join(column_or_null(columns, name) for name in names)
  cursor = source.execute(f"SELECT {selected} FROM snake_turns ORDER BY id")
  sql = f"INSERT INTO snake_turns ({','.join(names)}) VALUES ({','.join('?' for _ in names)})"
  count = 0
  while rows := cursor.fetchmany(batch_size):
    values = [tuple(row) for row in rows if row[0] in allowed]
    destination.executemany(sql, values)
    count += len(values)
  return count

def copy_replays(source, destination, allowed:set[str], batch_size:int) -> dict[str, int]:
  return {
    "game_snakes": copy_game_snakes(source, destination, allowed, batch_size),
    "turns": copy_turns(source, destination, allowed, batch_size),
    "snake_turns": copy_snake_turns(source, destination, allowed, batch_size),
  }

def verify(destination:sqlite3.Connection, expected:dict[str, int]) -> None:
  for table, count in expected.items():
    actual = int(destination.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0])
    if actual != count:
      raise RuntimeError(f"{table} count mismatch: expected {count}, got {actual}")
  foreign_keys = destination.execute("PRAGMA foreign_key_check").fetchall()
  if foreign_keys:
    raise RuntimeError(f"Foreign-key verification failed with {len(foreign_keys)} errors")
  integrity = destination.execute("PRAGMA integrity_check").fetchone()[0]
  if integrity != "ok":
    raise RuntimeError(f"Integrity check failed: {integrity}")

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

def main() -> None:
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("--base", required=True, type=Path, help="Cleaned database")
  parser.add_argument("--delta", required=True, type=Path, help="Database written during cleanup")
  parser.add_argument("--destination", required=True, type=Path)
  parser.add_argument("--minimum-quality", choices=("low", "medium", "high"), default="medium")
  parser.add_argument("--batch-size", type=int, default=10_000)
  parser.add_argument("--busy-timeout-ms", type=int, default=60_000)
  parser.add_argument("--replace-base", action="store_true", help="Replace base after verification and keep a timestamped backup")
  args = parser.parse_args()

  base_path = args.base.expanduser().resolve()
  delta_path = args.delta.expanduser().resolve()
  destination_path = args.destination.expanduser().resolve()
  if len({base_path, delta_path, destination_path}) != 3:
    raise SystemExit("Base, delta, and destination must be different paths")

  started = perf_counter()
  base = open_source(base_path)
  delta = open_source(delta_path)
  destination = create_destination(destination_path, max(1000, args.busy_timeout_ms))
  try:
    base_games = read_games(base)
    delta_games = read_games(delta)
    duplicates = duplicate_ids(base_games, delta_games)
    delta_quality = analyze_quality(delta)

    base_count, base_replays, _ = insert_games(destination, base_games, set())
    delta_count, delta_replays, tiers = insert_games(
      destination, delta_games, duplicates, delta_quality, args.minimum_quality,
    )
    base_rows = copy_replays(base, destination, base_replays, max(1, args.batch_size))
    delta_rows = copy_replays(delta, destination, delta_replays - duplicates, max(1, args.batch_size))
    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")
    expected = {
      "games": base_count + delta_count,
      "game_snakes": base_rows["game_snakes"] + delta_rows["game_snakes"],
      "turns": base_rows["turns"] + delta_rows["turns"],
      "snake_turns": base_rows["snake_turns"] + delta_rows["snake_turns"],
    }
    verify(destination, expected)
  except Exception:
    destination.close()
    base.close()
    delta.close()
    for suffix in ("", "-wal", "-shm"):
      Path(f"{destination_path}{suffix}").unlink(missing_ok=True)
    raise
  destination.close()
  base.close()
  delta.close()

  print(f"merged: base games={base_count:,}, delta games={delta_count:,}, identical duplicates skipped={len(duplicates):,}")
  print(f"delta quality: {dict(tiers)}")
  print(f"verified: {expected}; elapsed: {perf_counter() - started:.1f}s")
  if args.replace_base:
    backup = replace_base(base_path, destination_path)
    print(f"base replaced; backup kept at: {backup}")
  else:
    print(f"merged database created at: {destination_path}")

if __name__ == "__main__":
  main()