"""PostgreSQL gameplay backend using asyncpg. JSON columns use the JSONB type so PostgreSQL stores them in a binary, decomposed format and automatically compresses large values via TOAST (Oversized-Attribute Storage Technique). No application-level serialisation/deserialisation round-trip is needed for reads — asyncpg decodes JSONB rows directly into Python dicts/lists. 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 asyncio, json, logging, sqlite3, sys from datetime import datetime, timezone from pathlib import Path from urllib.parse import urlparse, urlunparse from .Template import GameplayBackendTemplate 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 # DDL --------------------------------------------------------------------- # _DDL = """ CREATE TABLE IF NOT EXISTS games ( game_id TEXT PRIMARY KEY, started_at TIMESTAMPTZ NOT NULL, ended_at TIMESTAMPTZ, width INTEGER, height INTEGER, source TEXT, map_name TEXT, ruleset_name TEXT, ruleset_version TEXT, your_snake_id TEXT, your_snake_name TEXT, your_snake_type TEXT, your_snake_version TEXT, game_type TEXT, winner_name TEXT, winner_you BOOLEAN NOT NULL DEFAULT FALSE, final_turn INTEGER NOT NULL DEFAULT 0, status TEXT NOT NULL DEFAULT 'running' ); CREATE TABLE IF NOT EXISTS turns ( id BIGSERIAL PRIMARY KEY, game_id TEXT NOT NULL REFERENCES games(game_id) ON DELETE CASCADE, turn INTEGER NOT NULL, observed_at TIMESTAMPTZ NOT NULL, my_move TEXT, my_thinking JSONB, board_state JSONB NOT NULL, snakes JSONB NOT NULL, you JSONB NOT NULL, food JSONB NOT NULL, hazards JSONB NOT NULL, UNIQUE (game_id, turn) ); CREATE TABLE IF NOT EXISTS snake_turns ( id BIGSERIAL PRIMARY KEY, game_id TEXT NOT NULL REFERENCES games(game_id) ON DELETE CASCADE, turn INTEGER NOT NULL, snake_id TEXT NOT NULL, snake_name TEXT, health INTEGER, length INTEGER, head_x INTEGER, head_y INTEGER, body JSONB NOT NULL, is_you BOOLEAN NOT NULL DEFAULT FALSE, inferred_move TEXT, latency TEXT, UNIQUE (game_id, turn, snake_id) ); 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); """ # Schema evolution: add new columns to existing tables (idempotent). _ALTER_DDL = """ ALTER TABLE games ADD COLUMN IF NOT EXISTS game_type TEXT; ALTER TABLE games ADD COLUMN IF NOT EXISTS your_snake_type TEXT; ALTER TABLE games ADD COLUMN IF NOT EXISTS your_snake_version TEXT; ALTER TABLE games ADD COLUMN IF NOT EXISTS winner_name TEXT; ALTER TABLE turns ADD COLUMN IF NOT EXISTS my_thinking JSONB; ALTER TABLE snake_turns ADD COLUMN IF NOT EXISTS latency TEXT; """ # Force TOAST compression on the large JSONB columns so that even # moderately-sized payloads get compressed on-disk. _TOAST_DDL = """ ALTER TABLE turns ALTER COLUMN board_state SET STORAGE EXTENDED; ALTER TABLE turns ALTER COLUMN snakes SET STORAGE EXTENDED; ALTER TABLE turns ALTER COLUMN you SET STORAGE EXTENDED; ALTER TABLE turns ALTER COLUMN food SET STORAGE EXTENDED; ALTER TABLE turns ALTER COLUMN hazards SET STORAGE EXTENDED; ALTER TABLE snake_turns ALTER COLUMN body SET STORAGE EXTENDED; """ 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 # ── 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. A DSN has no database name when its path component is empty or ``/``. In that case ``battlesnake`` is appended so asyncpg gets a complete connection string without the caller having to remember to add one. """ parsed = urlparse(dsn) db = parsed.path.lstrip("/") if db: return dsn new_path = f"/{cls._DEFAULT_DB_NAME}" return urlunparse(parsed._replace(path=new_path)) # ── pool / schema ────────────────────────────────────────────────────────── async def initialize(self) -> None: """Eagerly create the connection pool on startup so schema init and SQLite migration run immediately rather than on the first game request.""" 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_migrate_from_sqlite() return self._pool async def _ensure_database_exists(self, asyncpg, target_dsn:str) -> None: """Connect to the postgres maintenance DB and CREATE the target database if it does not already exist. Uses a plain connection (not a pool) so the CREATE DATABASE statement can run outside any transaction.""" 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: # Fall back to connecting without specifying a database — some setups # (e.g. Cloud SQL, managed PG) 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: await conn.execute(_DDL) await conn.execute(_ALTER_DDL) # TOAST storage hints are idempotent; ignore errors on repeated runs. try: await conn.execute(_TOAST_DDL) except Exception as exc: logger.debug(f"TOAST DDL skipped (likely already set): {exc}") # ── sqlite migration ─────────────────────────────────────────────────────── async def _maybe_migrate_from_sqlite(self) -> None: if not self._sqlite_migration_path: return src = Path(self._sqlite_migration_path) if not src.exists(): return logger.info(f"SQLite migration: found {src}, starting migration to PostgreSQL …") try: games, turns, snake_turns = await asyncio.to_thread(self._read_sqlite_data_sync, str(src)) await self._insert_migrated_data(games, turns, snake_turns) done_path = src.with_suffix(".migrated") src.rename(done_path) logger.info( f"SQLite migration complete: {len(games)} games, {len(turns)} turns, " f"{len(snake_turns)} snake_turns migrated. " f"Source file renamed to {done_path.name}" ) except Exception: logger.exception("SQLite migration failed — PostgreSQL data is untouched, original SQLite file kept") def _read_sqlite_data_sync(self, db_path:str) -> tuple[list[sqlite3.Row], list[sqlite3.Row], list[sqlite3.Row]]: conn = sqlite3.connect(db_path, timeout=30, isolation_level=None) conn.row_factory = sqlite3.Row try: games = conn.execute(""" 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_names_json, winner_you, final_turn, status FROM games ORDER BY started_at ASC """).fetchall() turns = conn.execute(""" 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 game_id ASC, turn ASC """).fetchall() snake_turns = conn.execute(""" 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 game_id ASC, turn ASC, snake_id ASC """).fetchall() finally: conn.close() return games, turns, snake_turns def _parse_ts(self, value:str|None) -> datetime|None: """Parse an ISO-8601 TEXT timestamp from SQLite into a timezone-aware datetime.""" ts = self._parse_utc_timestamp(value) return ts # already UTC-aware from base class helper def _parse_json(self, value: str|None) -> object: if not value: return None try: return json.loads(value) except (json.JSONDecodeError, TypeError): return None async def _insert_migrated_data(self, games:list, turns:list, snake_turns:list) -> None: assert self._pool is not None async with self._pool.acquire() as conn: async with conn.transaction(): # games ───────────────────────────────────────────────────────────── # winner_name is TEXT — no cast needed. await conn.executemany(""" 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 ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18) ON CONFLICT (game_id) DO NOTHING """, [ ( row["game_id"], self._parse_ts(row["started_at"]), self._parse_ts(row["ended_at"]), row["width"], row["height"], row["source"], row["map_name"], row["ruleset_name"], row["ruleset_version"], row["your_snake_id"], row["your_snake_name"], row["your_snake_type"], row["your_snake_version"], row["game_type"], (self._parse_json(row["winner_names_json"]) or [None])[0], bool(row["winner_you"]), row["final_turn"], row["status"], ) for row in games ], ) await conn.executemany(""" INSERT INTO turns ( game_id, turn, observed_at, my_move, my_thinking, board_state, snakes, you, food, hazards ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) ON CONFLICT (game_id, turn) DO NOTHING """, [ ( row["game_id"], row["turn"], self._parse_ts(row["observed_at"]), row["my_move"], self._parse_json(row["my_thinking_json"]), self._parse_json(row["board_state_json"]), self._parse_json(row["snakes_json"]), self._parse_json(row["you_json"]), self._parse_json(row["food_json"]), self._parse_json(row["hazards_json"]), ) for row in turns ], ) # snake_turns await conn.executemany(""" INSERT INTO snake_turns ( game_id, turn, snake_id, snake_name, health, length, head_x, head_y, body, is_you, inferred_move, latency ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12) ON CONFLICT (game_id, turn, snake_id) DO NOTHING """, [ ( row["game_id"], row["turn"], row["snake_id"], row["snake_name"], row["health"], row["length"], row["head_x"], row["head_y"], self._parse_json(row["body_json"]), bool(row["is_you"]), row["inferred_move"], row["latency"], ) for row in snake_turns ], ) # ── helpers ──────────────────────────────────────────────────────────────── def _utc_now_ts(self) -> datetime: return datetime.now(timezone.utc) # The pool init callback registers JSON/JSONB codecs so asyncpg automatically # encodes Python dicts/lists on write and decodes them on read. # ── 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 = game_state.get("game", {}) board = game_state.get("board", {}) snakes = self._extract_snakes(game_state) you = self._extract_you(game_state) game_id = game.get("id") turn = int(game_state.get("turn", 0)) pool = await self._get_pool() async with pool.acquire() as conn: async with conn.transaction(): await conn.execute(""" INSERT INTO turns ( game_id, turn, observed_at, my_move, my_thinking, board_state, snakes, you, food, hazards ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) ON CONFLICT (game_id, turn) DO UPDATE SET observed_at = EXCLUDED.observed_at, my_move = EXCLUDED.my_move, my_thinking = EXCLUDED.my_thinking, board_state = EXCLUDED.board_state, snakes = EXCLUDED.snakes, you = EXCLUDED.you, food = EXCLUDED.food, hazards = EXCLUDED.hazards """, game_id, turn, self._utc_now_ts(), my_move, my_thinking, board, snakes, you, board.get("food", []), board.get("hazards", []), ) previous_positions:dict[str, tuple[int, int]] = {} if turn > 0: prev_rows = await conn.fetch(""" SELECT snake_id, head_x, head_y FROM snake_turns WHERE game_id = $1 AND turn = $2 """, game_id, turn - 1, ) previous_positions = { row["snake_id"]: (int(row["head_x"]), int(row["head_y"])) for row in prev_rows if row["head_x"] is not None and row["head_y"] is not None } you_id = you.get("id") for params in self._build_snake_turn_params(snakes, you_id, game_id, turn, previous_positions): (p_game_id, p_turn, p_snake_id, p_name, p_health, p_length, p_head_x, p_head_y, p_body, p_is_you, p_inferred, p_latency) = params await conn.execute(""" INSERT INTO snake_turns ( game_id, turn, snake_id, snake_name, health, length, head_x, head_y, body, is_you, inferred_move, latency ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12) ON CONFLICT (game_id, turn, snake_id) DO UPDATE SET snake_name = EXCLUDED.snake_name, health = EXCLUDED.health, length = EXCLUDED.length, head_x = EXCLUDED.head_x, head_y = EXCLUDED.head_y, body = EXCLUDED.body, is_you = EXCLUDED.is_you, inferred_move = EXCLUDED.inferred_move, latency = EXCLUDED.latency """, p_game_id, p_turn, p_snake_id, p_name, p_health, p_length, p_head_x, p_head_y, p_body, p_is_you, p_inferred, p_latency, ) await conn.execute(""" UPDATE games SET final_turn = GREATEST(final_turn, $1) WHERE game_id = $2 """, turn, game_id, ) 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, ) # ── 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 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"] your_snake_id = row["your_snake_id"] final_turn = int(row["final_turn"] or 0) snake_rows = await conn.fetch(""" SELECT snake_id, snake_name FROM snake_turns WHERE game_id = $1 AND turn = $2 ORDER BY is_you DESC, snake_name ASC """, game_id, final_turn, ) if len(snake_rows) == 0: latest_row = await conn.fetchrow( "SELECT MAX(turn) AS latest_turn FROM snake_turns WHERE game_id = $1", game_id, ) if latest_row is not None and latest_row["latest_turn"] is not None: final_turn = int(latest_row["latest_turn"]) snake_rows = await conn.fetch(""" SELECT snake_id, snake_name FROM snake_turns WHERE game_id = $1 AND turn = $2 ORDER BY is_you DESC, snake_name ASC """, game_id, final_turn, ) winner_you, survivor_name = self._calculate_survivor(snake_rows, 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 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 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 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(""" SELECT 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 FROM games WHERE game_id = $1 """, game_id, ) if game_row is None: return None turn_rows = await conn.fetch(""" SELECT turn, observed_at, my_move, my_thinking AS my_thinking_json, board_state AS board_state_json, food AS food_json, hazards AS hazards_json, you AS you_json FROM turns WHERE game_id = $1 ORDER BY turn ASC """, game_id, ) snake_rows = await conn.fetch(""" SELECT turn, snake_id, snake_name, health, length, head_x, head_y, body AS body_json, is_you, inferred_move, latency FROM snake_turns WHERE game_id = $1 ORDER BY turn ASC, is_you DESC, snake_name ASC """, game_id, ) return self._build_game_replay_output(game_row, turn_rows, snake_rows, decode_json=lambda x: x) # ── lifecycle ────────────────────────────────────────────────────────────── async def close(self) -> None: if self._pool is not None: await self._pool.close() self._pool = None