snake-python/scripts/migrate_sqlite_to_postgresql.py

#!/usr/bin/env python3
"""Stream a normalized Battlesnake SQLite database into PostgreSQL.

The source is opened read-only. Rows are copied in bounded batches through
temporary PostgreSQL tables, then inserted idempotently with ON CONFLICT. The
script verifies source/inserted row counts and never modifies the SQLite file.
"""

from __future__ import annotations

import argparse
import asyncio
from datetime import datetime
import json
from pathlib import Path
import sqlite3
from time import perf_counter

from server.database.backend.PostgresqlGameplayBackend import PostgresqlGameplayBackend

TABLES = (
  (
    "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",
    ),
    "game_id",
  ),
  (
    "game_snakes",
    ("game_id", "snake_id", "snake_name", "is_you", "customizations"),
    "game_id, snake_id",
  ),
  (
    "turns",
    (
      "game_id", "turn", "observed_at", "my_move", "my_thinking",
      "board_state", "snakes", "you", "food", "hazards",
    ),
    "game_id, turn",
  ),
  (
    "snake_turns",
    (
      "game_id", "turn", "snake_id", "snake_name", "health", "length",
      "head_x", "head_y", "body", "is_you", "inferred_move", "latency",
    ),
    "game_id, turn, snake_id",
  ),
)

SQLITE_SELECTS = {
  "games": """
    SELECT 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
    FROM games ORDER BY game_id
  """,
  "game_snakes": """
    SELECT game_id, snake_id, snake_name, is_you, customizations_json
    FROM game_snakes ORDER BY game_id, snake_id
  """,
  "turns": """
    SELECT game_id, turn, observed_at, my_move, my_thinking_json,
      board_state_json, snakes_json, you_json, food_json, hazards_json
    FROM turns ORDER BY id
  """,
  "snake_turns": """
    SELECT game_id, turn, snake_id, snake_name, health, length, head_x, head_y,
      body_json, is_you, inferred_move, latency
    FROM snake_turns ORDER BY id
  """,
}

TIMESTAMP_FIELDS = {"started_at", "ended_at", "observed_at"}
BOOLEAN_FIELDS = {"winner_you", "has_replay", "is_you"}
JSON_FIELDS = {
  "quality_reasons", "customizations", "my_thinking", "board_state",
  "snakes", "you", "food", "hazards", "body",
}
JSON_DEFAULTS = {
  "customizations": {}, "board_state": {}, "snakes": [], "you": {},
  "food": [], "hazards": [], "body": [],
}

def parse_timestamp(value: str | None) -> datetime | None:
  if not value:
    return None
  parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
  return parsed

def parse_json(value: str | None, default: object = None) -> object:
  if value in (None, ""):
    return default
  try:
    return json.loads(value)
  except (json.JSONDecodeError, TypeError) as exc:
    raise ValueError(f"Invalid JSON value: {str(value)[:120]}") from exc

def transform_row(columns: tuple[str, ...], row: sqlite3.Row) -> tuple:
  output = []
  for column, value in zip(columns, row, strict=True):
    if column in TIMESTAMP_FIELDS:
      value = parse_timestamp(value)
    elif column in BOOLEAN_FIELDS:
      value = bool(value)
    elif column in JSON_FIELDS:
      parsed = parse_json(value, JSON_DEFAULTS.get(column))
      value = None if parsed is None else json.dumps(parsed, separators=(",", ":"))
    output.append(value)
  return tuple(output)

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

async def import_table(pool, source: sqlite3.Connection, table: str,
                       columns: tuple[str, ...], conflict_columns: str,
                       batch_size: int) -> tuple[int, int]:
  source_count = int(source.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0])
  async with pool.acquire() as connection:
    before = int(await connection.fetchval(f"SELECT COUNT(*) FROM {table}"))
    stage = f"migration_{table}"
    await connection.execute(
      f"CREATE TEMP TABLE {stage} (LIKE {table} INCLUDING DEFAULTS)"
    )
    json_columns = [column for column in columns if column in JSON_FIELDS]
    for column in json_columns:
      await connection.execute(
        f"ALTER TABLE {stage} ALTER COLUMN {column} TYPE TEXT USING {column}::text"
      )
    cursor = source.execute(SQLITE_SELECTS[table])
    copied = 0
    while rows := cursor.fetchmany(batch_size):
      records = [transform_row(columns, row) for row in rows]
      async with connection.transaction():
        await connection.copy_records_to_table(stage, records=records, columns=columns)
        selected = ", ".join(columns)
        source_expressions = ", ".join(
          f"{column}::jsonb" if column in JSON_FIELDS else column
          for column in columns
        )
        await connection.execute(
          f"INSERT INTO {table} ({selected}) SELECT {source_expressions} FROM {stage} "
          f"ON CONFLICT ({conflict_columns}) DO NOTHING"
        )
        await connection.execute(f"TRUNCATE {stage}")
      copied += len(records)
      print(f"{table}: streamed {copied:,}/{source_count:,}", flush=True)
    after = int(await connection.fetchval(f"SELECT COUNT(*) FROM {table}"))
  inserted = after - before
  if copied != source_count:
    raise RuntimeError(f"{table}: source changed while reading ({source_count} -> {copied})")
  print(f"{table}: source={source_count:,}, inserted={inserted:,}, conflicts={source_count - inserted:,}")
  return source_count, inserted

async def migrate(source_path: Path, dsn: str, batch_size: int) -> None:
  source = open_source(source_path)
  try:
    quick_check = source.execute("PRAGMA quick_check").fetchone()[0]
    if quick_check != "ok":
      raise RuntimeError(f"SQLite quick_check failed: {quick_check}")

    backend = PostgresqlGameplayBackend(dsn=dsn)
    await backend.initialize()
    pool = await backend._get_pool()
    started = perf_counter()
    results = {}
    try:
      for table, columns, conflicts in TABLES:
        results[table] = await import_table(
          pool, source, table, columns, conflicts, max(100, batch_size),
        )
      async with pool.acquire() as connection:
        invalid = int(await connection.fetchval("""
          SELECT
            (SELECT COUNT(*) FROM turns t LEFT JOIN games g USING (game_id) WHERE g.game_id IS NULL) +
            (SELECT COUNT(*) FROM game_snakes s LEFT JOIN games g USING (game_id) WHERE g.game_id IS NULL) +
            (SELECT COUNT(*) FROM snake_turns s LEFT JOIN games g USING (game_id) WHERE g.game_id IS NULL)
        """))
        if invalid:
          raise RuntimeError(f"PostgreSQL foreign-key verification found {invalid} orphan rows")
        counts = {
          table: int(await connection.fetchval(f"SELECT COUNT(*) FROM {table}"))
          for table, _, _ in TABLES
        }
      print(f"verified PostgreSQL counts: {counts}")
      print(f"migration elapsed: {perf_counter() - started:.1f}s")
    finally:
      await backend.close()
  finally:
    source.close()

def main() -> None:
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("--source", required=True, type=Path)
  parser.add_argument("--dsn", required=True)
  parser.add_argument("--batch-size", type=int, default=10_000)
  args = parser.parse_args()
  asyncio.run(migrate(args.source.expanduser().resolve(), args.dsn, args.batch_size))

if __name__ == "__main__":
  main()