snake-python/server/database/backend/PostgresqlGameplayBackend.py

"""PostgreSQL gameplay backend using asyncpg.

Running games are stored as compact rows (``game_turns``/``snake_states``).
When a game ends it is packed into one zstd blob in ``game_replays`` and its
rows are removed. See ``server.database.compact`` for the encodings.

Connection: pass a DSN via the `dsn` constructor argument, e.g.
  postgresql://user:password@host:5432/dbname

or set GAMEPLAY_DB_PG_DSN in the environment.
"""

import json, logging, sys
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import urlparse, urlunparse

from .Template import GameplayBackendTemplate
from .postgresql_schema import DROP_REDUNDANT_INDEXES, TABLES_DDL
from server.database.compact.cache import RunningGame, RunningGameCache
from server.database.compact.codec import (
  decode_body, decode_move, decode_points, encode_body, encode_move, encode_points,
)
from server.database.compact.games import GAME_COLUMNS, normalize_game
from server.database.compact.records import (
  GameBundle, SnakeIdentity, SnakeState, TurnRecord, merge_records, records_from_game_state,
)
from server.database.compact.replay_blob import REPLAY_FORMAT, pack_replay, unpack_replay

logger = logging.getLogger(__name__)
if not logger.handlers:
  _handler = logging.StreamHandler(stream=sys.stdout)
  _handler.setFormatter(logging.Formatter(fmt="%(levelname)s %(module)s: %(message)s"))
  logger.addHandler(_handler)
  logger.propagate = False

_SCHEMA_LOCK_ID = 7_315_002

_GAME_SELECT = """
  SELECT gid, game_id, started_at, ended_at, width, height, source, map_name,
    ruleset_name, ruleset_version, game_type, your_snake_id, your_snake_name,
    your_snake_type, your_snake_version, winner_name, winner_you, final_turn,
    status, has_replay, quality_status, quality_score, quality_tier, quality_reasons
  FROM games
"""

_UPSERT_SNAKES = """
  INSERT INTO game_snakes (game_id, snake_id, snake_name, is_you, customizations)
  SELECT $1, s.snake_id, s.snake_name, s.is_you, s.customizations::jsonb
  FROM unnest($2::text[], $3::text[], $4::bool[], $5::text[])
    AS s(snake_id, snake_name, is_you, customizations)
  ON CONFLICT (game_id, snake_id) DO UPDATE SET
    snake_name = EXCLUDED.snake_name,
    is_you = EXCLUDED.is_you,
    customizations = EXCLUDED.customizations
  RETURNING snake_id, gsid
"""

_UPSERT_TURN = """
  INSERT INTO game_turns (gid, turn, observed_at, my_move, thinking, food, hazards)
  VALUES ($1, $2, $3, $4, $5, $6, $7)
  ON CONFLICT (gid, turn) DO UPDATE SET
    observed_at = EXCLUDED.observed_at,
    my_move = EXCLUDED.my_move,
    thinking = EXCLUDED.thinking,
    food = EXCLUDED.food,
    hazards = EXCLUDED.hazards
"""

_UPSERT_STATES = """
  INSERT INTO snake_states (gsid, turn, health, latency, body)
  SELECT * FROM unnest($1::int[], $2::int[], $3::smallint[], $4::smallint[], $5::bytea[])
  ON CONFLICT (gsid, turn) DO UPDATE SET
    health = EXCLUDED.health,
    latency = EXCLUDED.latency,
    body = EXCLUDED.body
"""

class PostgresqlGameplayBackend(GameplayBackendTemplate):
  """Async PostgreSQL backend.  A connection pool is created lazily on the
  first method call and reused for the lifetime of the object.

  Requires: pip install asyncpg
  """

  def __init__(self, dsn:str, min_size:int=1, max_size:int=5, sqlite_migration_path:str|None=None):
    self._dsn = dsn
    self._min_size = min_size
    self._max_size = max_size
    self._sqlite_migration_path = sqlite_migration_path
    self._pool = None  # asyncpg.Pool, typed at runtime
    self._running = RunningGameCache()

  # ── DSN normalisation ──────────────────────────────────────────────────────

  _DEFAULT_DB_NAME = "battlesnake"

  @classmethod
  def _ensure_db_name(cls, dsn:str) -> str:
    """Return *dsn* with a database name appended when none is present."""
    parsed = urlparse(dsn)
    if parsed.path.lstrip("/"):
      return dsn
    return urlunparse(parsed._replace(path=f"/{cls._DEFAULT_DB_NAME}"))

  # ── pool / schema ──────────────────────────────────────────────────────────

  async def initialize(self) -> None:
    """Eagerly create the pool so schema init and legacy SQLite import run on startup."""
    await self._get_pool()

  async def _get_pool(self):
    if self._pool is None:
      try:
        import asyncpg  # noqa: PLC0415
      except ImportError as exc:
        raise ImportError(
          "asyncpg is required for the PostgreSQL gameplay backend. "
          "Install it with: pip install asyncpg"
        ) from exc

      target_dsn = self._ensure_db_name(self._dsn)
      await self._ensure_database_exists(asyncpg, target_dsn)

      async def _init_conn(conn) -> None:
        await conn.set_type_codec('jsonb', encoder=json.dumps, decoder=json.loads, schema='pg_catalog')
        await conn.set_type_codec('json',  encoder=json.dumps, decoder=json.loads, schema='pg_catalog')

      self._pool = await asyncpg.create_pool(
        dsn=target_dsn,
        min_size=self._min_size,
        max_size=self._max_size,
        init=_init_conn,
      )
      await self._initialize_schema()
      await self._maybe_import_legacy_sqlite()
    return self._pool

  async def _ensure_database_exists(self, asyncpg, target_dsn:str) -> None:
    """CREATE the target database through the maintenance DB when missing."""
    parsed = urlparse(target_dsn)
    db_name = parsed.path.lstrip("/")
    maintenance_dsn = urlunparse(parsed._replace(path="/postgres"))
    try:
      conn = await asyncpg.connect(dsn=maintenance_dsn)
    except Exception:
      # Some managed setups disallow direct access to 'postgres'.
      maintenance_dsn = urlunparse(parsed._replace(path=""))
      conn = await asyncpg.connect(dsn=maintenance_dsn)
    try:
      exists = await conn.fetchval("SELECT 1 FROM pg_database WHERE datname = $1", db_name)
      if not exists:
        await conn.execute(f'CREATE DATABASE "{db_name}"')
        logger.info(f"PostgreSQL: created database '{db_name}'")
    finally:
      await conn.close()

  async def _initialize_schema(self) -> None:
    assert self._pool is not None
    async with self._pool.acquire() as conn:
      # Several workers may start at once; serialize DDL.
      await conn.execute("SELECT pg_advisory_lock($1)", _SCHEMA_LOCK_ID)
      try:
        await conn.execute(TABLES_DDL)
      finally:
        await conn.execute("SELECT pg_advisory_unlock($1)", _SCHEMA_LOCK_ID)
      for statement in DROP_REDUNDANT_INDEXES:
        await conn.execute(statement)

  async def _maybe_import_legacy_sqlite(self) -> None:
    if not self._sqlite_migration_path:
      return
    source = Path(self._sqlite_migration_path)
    if not source.exists():
      return

    from server.database.legacy.migrate import migrate_legacy  # noqa: PLC0415
    from server.database.legacy.reader_sqlite import LegacySqliteReader  # noqa: PLC0415

    logger.info(f"Legacy SQLite import: found {source}, importing into PostgreSQL …")
    try:
      with LegacySqliteReader(source) as reader:
        stats = await migrate_legacy(reader, self)
      if stats.failed:
        logger.warning(f"Legacy SQLite import incomplete ({stats.failed} games failed); source kept")
        return
      done_path = source.with_suffix(".migrated")
      source.rename(done_path)
      logger.info(f"Legacy SQLite import complete: {stats.games} games. Source renamed to {done_path.name}")
    except Exception:
      logger.exception("Legacy SQLite import failed — original SQLite file kept")

  # ── helpers ────────────────────────────────────────────────────────────────

  def _utc_now_ts(self) -> datetime:
    return datetime.now(timezone.utc)

  async def _running_game(self, conn, game_id:str) -> RunningGame:
    game = self._running.get(game_id)
    if game is None:
      gid = await conn.fetchval("SELECT gid FROM games WHERE game_id = $1", game_id)
      if gid is None:
        raise LookupError(f"Turn for unknown game {game_id!r}")
      game = RunningGame(gid=gid)
    return game

  async def _upsert_identities(self, conn, game_id:str, identities:list[SnakeIdentity]) -> dict[str, int]:
    if not identities:
      return {}
    rows = await conn.fetch(
      _UPSERT_SNAKES, game_id,
      [identity.snake_id for identity in identities],
      [identity.name for identity in identities],
      [identity.is_you for identity in identities],
      [json.dumps(identity.customizations or {}) for identity in identities],
    )
    return {row["snake_id"]: row["gsid"] for row in rows}

  async def _write_rows(self, conn, gid:int, gsids:dict[str, int], turns:list[TurnRecord], states:list[SnakeState]) -> None:
    if turns:
      await conn.executemany(_UPSERT_TURN, [
        (
          gid, record.turn, record.observed_at or self._utc_now_ts(), encode_move(record.move),
          record.thinking, encode_points(record.food), encode_points(record.hazards),
        )
        for record in turns
      ])
    states = [state for state in states if state.snake_id in gsids]
    if states:
      await conn.execute(
        _UPSERT_STATES,
        [gsids[state.snake_id] for state in states],
        [state.turn for state in states],
        [state.health for state in states],
        [state.latency for state in states],
        [encode_body(state.body) for state in states],
      )

  async def _load_records(self, conn, game_id:str, gid:int) -> tuple[list[SnakeIdentity], list[TurnRecord], list[SnakeState]]:
    identity_rows = await conn.fetch("""
      SELECT snake_id, snake_name, is_you, customizations
      FROM game_snakes WHERE game_id = $1
      """, game_id)
    identities = [
      SnakeIdentity(row["snake_id"], row["snake_name"], bool(row["is_you"]), row["customizations"] or {})
      for row in identity_rows
    ]
    blob = await conn.fetchrow("SELECT format, data FROM game_replays WHERE gid = $1", gid)
    turn_rows = await conn.fetch("""
      SELECT turn, observed_at, my_move, thinking, food, hazards
      FROM game_turns WHERE gid = $1 ORDER BY turn
      """, gid)
    state_rows = await conn.fetch("""
      SELECT gs.snake_id, s.turn, s.health, s.latency, s.body
      FROM game_snakes AS gs
      JOIN snake_states AS s ON s.gsid = gs.gsid
      WHERE gs.game_id = $1
      """, game_id)
    turns = [
      TurnRecord(
        turn=row["turn"], observed_at=row["observed_at"], move=decode_move(row["my_move"]),
        thinking=row["thinking"], food=decode_points(row["food"]), hazards=decode_points(row["hazards"]),
      )
      for row in turn_rows
    ]
    states = [
      SnakeState(row["snake_id"], row["turn"], row["health"], row["latency"], decode_body(row["body"]))
      for row in state_rows
    ]
    if blob is not None:
      packed_turns, packed_states = unpack_replay(blob["data"], blob["format"])
      turns, states = merge_records(packed_turns, packed_states, turns, states)
    return identities, turns, states

  async def _store_replay(self, conn, game_id:str, gid:int, turns:list[TurnRecord], states:list[SnakeState]) -> None:
    await conn.execute("""
      INSERT INTO game_replays (gid, format, turn_count, data) VALUES ($1, $2, $3, $4)
      ON CONFLICT (gid) DO UPDATE SET
        format = EXCLUDED.format, turn_count = EXCLUDED.turn_count, data = EXCLUDED.data
      """, gid, REPLAY_FORMAT, len(turns), pack_replay(turns, states))
    await conn.execute("DELETE FROM game_turns WHERE gid = $1", gid)
    await conn.execute("""
      DELETE FROM snake_states
      WHERE gsid IN (SELECT gsid FROM game_snakes WHERE game_id = $1)
      """, game_id)

  async def _pack_game(self, conn, game_id:str) -> None:
    """Move a game's rows into its replay blob (idempotent)."""
    async with conn.transaction():
      gid = await conn.fetchval("SELECT gid FROM games WHERE game_id = $1 FOR UPDATE", game_id)
      if gid is None:
        return
      has_rows = await conn.fetchval("SELECT EXISTS (SELECT 1 FROM game_turns WHERE gid = $1)", gid)
      if not has_rows:
        return
      _, turns, states = await self._load_records(conn, game_id, gid)
      await self._store_replay(conn, game_id, gid, turns, states)

  # ── write methods ──────────────────────────────────────────────────────────

  async def record_game_start(self, game_state:dict, snake_type:str|None=None, snake_version:str|None=None) -> None:
    game = game_state.get("game", {})
    board = game_state.get("board", {})
    you = self._extract_you(game_state)
    ruleset = game.get("ruleset", {})
    game_type = self._derive_game_type(board, ruleset)

    pool = await self._get_pool()
    async with pool.acquire() as conn:
      await conn.execute("""
        INSERT INTO games (
          game_id, started_at, width, height, source, map_name,
          ruleset_name, ruleset_version, your_snake_id, your_snake_name,
          your_snake_type, your_snake_version, game_type, status
        ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,'running')
        ON CONFLICT (game_id) DO UPDATE SET
          width             = EXCLUDED.width,
          height            = EXCLUDED.height,
          source            = EXCLUDED.source,
          map_name          = EXCLUDED.map_name,
          ruleset_name      = EXCLUDED.ruleset_name,
          ruleset_version   = EXCLUDED.ruleset_version,
          your_snake_id     = EXCLUDED.your_snake_id,
          your_snake_name   = EXCLUDED.your_snake_name,
          your_snake_type   = EXCLUDED.your_snake_type,
          your_snake_version = EXCLUDED.your_snake_version,
          game_type         = EXCLUDED.game_type,
          status            = 'running'
        """,
        game.get("id"),
        self._utc_now_ts(),
        board.get("width"),
        board.get("height"),
        game.get("source"),
        game.get("map"),
        ruleset.get("name"),
        ruleset.get("version"),
        you.get("id"),
        you.get("name"),
        snake_type,
        snake_version,
        game_type,
      )

  async def record_turn(self, game_state:dict, my_move:str|None, my_thinking:dict|None=None) -> None:
    game_id = game_state.get("game", {}).get("id")
    pool = await self._get_pool()
    async with pool.acquire() as conn:
      async with conn.transaction():
        running = await self._running_game(conn, game_id)
        record, states, identities, points = records_from_game_state(
          game_state, my_move, my_thinking, running.points,
        )
        missing = [identity for identity in identities if identity.snake_id not in running.snake_ids]
        if missing:
          running.snake_ids.update(await self._upsert_identities(conn, game_id, identities))
        await self._write_rows(conn, running.gid, running.snake_ids, [record], states)
        await conn.execute(
          "UPDATE games SET final_turn = GREATEST(final_turn, $1) WHERE gid = $2",
          record.turn, running.gid,
        )
    running.points = points
    self._running.put(game_id, running)

  async def record_game_end(self, game_state:dict) -> None:
    game_id, winner_name, winner_you, turn = self._extract_game_end_params(game_state)

    pool = await self._get_pool()
    async with pool.acquire() as conn:
      await conn.execute("""
        UPDATE games
        SET ended_at    = $1,
            winner_name = $2,
            winner_you  = $3,
            final_turn  = GREATEST(final_turn, $4),
            status      = 'finished'
        WHERE game_id = $5
        """,
        self._utc_now_ts(), winner_name, winner_you, turn, game_id,
      )
      await self._pack_game(conn, game_id)
    self._running.drop(game_id)

  # ── stale game finalization ────────────────────────────────────────────────

  async def finalize_stale_running_games(self, stale_after_seconds:int=600) -> int:
    threshold = max(60, int(stale_after_seconds))
    now_utc = datetime.now(timezone.utc)
    finalized = 0

    pool = await self._get_pool()
    async with pool.acquire() as conn:
      rows = await conn.fetch("""
        SELECT gid, game_id, started_at, final_turn, your_snake_id
        FROM games
        WHERE status = 'running'
        ORDER BY started_at ASC
      """)

      for row in rows:
        started_at = row["started_at"]
        if started_at is None:
          continue
        if started_at.tzinfo is None:
          started_at = started_at.replace(tzinfo=timezone.utc)
        if (now_utc - started_at).total_seconds() < threshold:
          continue

        game_id = row["game_id"]
        identities, _, states = await self._load_records(conn, game_id, row["gid"])
        final_turn, survivors = self._survivors(identities, states, int(row["final_turn"] or 0))
        winner_you, survivor_name = self._calculate_survivor(survivors, row["your_snake_id"])

        tag = await conn.execute("""
          UPDATE games
          SET ended_at     = $1,
              winner_name  = $2,
              winner_you   = $3,
              final_turn   = GREATEST(final_turn, $4),
              status       = 'finished'
          WHERE game_id = $5 AND status = 'running'
          """,
          self._utc_now_ts(), survivor_name, winner_you, final_turn, game_id,
        )
        if tag and tag.endswith("1"):
          finalized += 1
          await self._pack_game(conn, game_id)
          self._running.drop(game_id)

      # Turns that arrived after a game was packed are folded into its blob.
      late = await conn.fetch("""
        SELECT DISTINCT g.game_id
        FROM game_turns AS t JOIN games AS g ON g.gid = t.gid
        WHERE g.status = 'finished'
      """)
      for late_row in late:
        await self._pack_game(conn, late_row["game_id"])

    return finalized

  # ── read methods ───────────────────────────────────────────────────────────

  async def get_summary(self, recent_limit:int=15) -> dict:
    pool = await self._get_pool()
    async with pool.acquire() as conn:
      totals = await conn.fetchrow("""
        SELECT
          COUNT(*)                                                       AS total_games,
          COUNT(*) FILTER (WHERE status = 'running')                    AS running_games,
          COUNT(*) FILTER (WHERE status = 'finished')                   AS finished_games,
          COUNT(*) FILTER (WHERE status = 'finished' AND winner_you)    AS wins,
          COUNT(*) FILTER (WHERE status = 'finished' AND NOT winner_you) AS losses,
          AVG(final_turn) FILTER (WHERE status = 'finished')            AS avg_turns
        FROM games
      """)

      by_type = await conn.fetch("""
        SELECT
          COALESCE(game_type, ruleset_name, 'unknown') AS type_label,
          COUNT(*)                                              AS total,
          COUNT(*) FILTER (WHERE winner_you)                   AS wins,
          COUNT(*) FILTER (WHERE NOT winner_you)               AS losses
        FROM games
        WHERE status = 'finished'
        GROUP BY type_label
        ORDER BY total DESC
      """)

      recent = await conn.fetch("""
        SELECT game_id, started_at, ended_at, map_name, ruleset_name, game_type,
          your_snake_name, your_snake_type, your_snake_version, winner_you, final_turn, status
        FROM games
        WHERE has_replay
        ORDER BY started_at DESC
        LIMIT $1
        """,
        max(1, int(recent_limit)),
      )

    return self._build_summary_output(totals, by_type, recent)

  async def list_games(self, limit:int=50) -> list[dict]:
    pool = await self._get_pool()
    async with pool.acquire() as conn:
      rows = await conn.fetch("""
        SELECT game_id, started_at, ended_at, map_name, source, ruleset_name, game_type,
          your_snake_name, your_snake_type, your_snake_version,
          winner_you, winner_name, final_turn, status
        FROM games
        WHERE has_replay
        ORDER BY started_at DESC
        LIMIT $1
        """,
        max(1, int(limit)),
      )

    return self._build_game_list_output(rows)

  async def get_game_replay(self, game_id:str) -> dict|None:
    pool = await self._get_pool()
    async with pool.acquire() as conn:
      game_row = await conn.fetchrow(_GAME_SELECT + " WHERE game_id = $1 AND has_replay", game_id)
      if game_row is None:
        return None
      identities, turns, states = await self._load_records(conn, game_id, game_row["gid"])

    return self._build_game_replay_output(game_row, identities, turns, states)

  # ── bundles (migration / merge) ────────────────────────────────────────────

  async def list_game_ids(self) -> list[str]:
    pool = await self._get_pool()
    async with pool.acquire() as conn:
      rows = await conn.fetch("SELECT game_id FROM games ORDER BY started_at, game_id")
    return [row["game_id"] for row in rows]

  async def read_game_bundle(self, game_id:str) -> GameBundle|None:
    pool = await self._get_pool()
    async with pool.acquire() as conn:
      game_row = await conn.fetchrow(_GAME_SELECT + " WHERE game_id = $1", game_id)
      if game_row is None:
        return None
      identities, turns, states = await self._load_records(conn, game_id, game_row["gid"])
    return GameBundle(normalize_game(dict(game_row)), identities, turns, states)

  async def import_game_bundle(self, bundle:GameBundle) -> None:
    game = normalize_game(bundle.game)
    columns = ", ".join(GAME_COLUMNS)
    placeholders = ", ".join(f"${index}" for index in range(1, len(GAME_COLUMNS) + 1))
    updates = ", ".join(f"{column} = EXCLUDED.{column}" for column in GAME_COLUMNS[1:])
    pool = await self._get_pool()
    async with pool.acquire() as conn:
      async with conn.transaction():
        gid = await conn.fetchval(
          f"INSERT INTO games ({columns}) VALUES ({placeholders}) "
          f"ON CONFLICT (game_id) DO UPDATE SET {updates} RETURNING gid",
          *[game[column] for column in GAME_COLUMNS],
        )
        gsids = await self._upsert_identities(conn, game["game_id"], bundle.snakes)
        if not bundle.turns or not game["has_replay"]:
          return
        if game["status"] == "finished":
          await self._store_replay(conn, game["game_id"], gid, bundle.turns, bundle.states)
        else:
          await self._write_rows(conn, gid, gsids, bundle.turns, bundle.states)
    self._running.drop(game["game_id"])

  # ── lifecycle ──────────────────────────────────────────────────────────────

  async def close(self) -> None:
    if self._pool is not None:
      await self._pool.close()
      self._pool = None