game-marathon/backup-service/restore.py

#!/usr/bin/env python3
"""
Restore PostgreSQL database from S3 backup.

Usage:
    python restore.py                       - List available backups
    python restore.py <filename>            - Restore from backup (cleans DB first)
    python restore.py <filename> --no-clean - Restore without cleaning DB first
"""
import gzip
import os
import subprocess
import sys

import boto3
from botocore.config import Config as BotoConfig
from botocore.exceptions import ClientError

from config import config


def create_s3_client():
    """Initialize S3 client."""
    return boto3.client(
        "s3",
        endpoint_url=config.S3_ENDPOINT_URL,
        aws_access_key_id=config.S3_ACCESS_KEY_ID,
        aws_secret_access_key=config.S3_SECRET_ACCESS_KEY,
        region_name=config.S3_REGION or "us-east-1",
        config=BotoConfig(signature_version="s3v4"),
    )


def list_backups(s3_client) -> list[tuple[str, float, str]]:
    """List all available backups."""
    print("Available backups:\n")

    try:
        paginator = s3_client.get_paginator("list_objects_v2")
        pages = paginator.paginate(
            Bucket=config.S3_BUCKET_NAME,
            Prefix=config.S3_BACKUP_PREFIX,
        )

        backups = []
        for page in pages:
            for obj in page.get("Contents", []):
                filename = obj["Key"].replace(config.S3_BACKUP_PREFIX, "")
                size_mb = obj["Size"] / (1024 * 1024)
                modified = obj["LastModified"].strftime("%Y-%m-%d %H:%M:%S")
                backups.append((filename, size_mb, modified))

        # Sort by date descending (newest first)
        backups.sort(key=lambda x: x[2], reverse=True)

        for filename, size_mb, modified in backups:
            print(f"  {filename} ({size_mb:.2f} MB) - {modified}")

        return backups

    except ClientError as e:
        print(f"Error listing backups: {e}")
        return []


def clean_database() -> None:
    """Drop and recreate public schema to clean the database."""
    print("Cleaning database (dropping and recreating public schema)...")

    env = os.environ.copy()
    env["PGPASSWORD"] = config.DB_PASSWORD

    # Drop and recreate public schema
    clean_sql = b"""
DROP SCHEMA public CASCADE;
CREATE SCHEMA public;
GRANT ALL ON SCHEMA public TO public;
"""

    cmd = [
        "psql",
        "-h",
        config.DB_HOST,
        "-p",
        config.DB_PORT,
        "-U",
        config.DB_USER,
        "-d",
        config.DB_NAME,
    ]

    result = subprocess.run(
        cmd,
        env=env,
        input=clean_sql,
        capture_output=True,
    )

    if result.returncode != 0:
        stderr = result.stderr.decode()
        if "ERROR" in stderr:
            raise Exception(f"Database cleanup failed: {stderr}")

    print("Database cleaned successfully!")


def restore_backup(s3_client, filename: str, clean_first: bool = True) -> None:
    """Download and restore backup."""
    key = f"{config.S3_BACKUP_PREFIX}{filename}"

    print(f"Downloading {filename} from S3...")
    try:
        response = s3_client.get_object(
            Bucket=config.S3_BUCKET_NAME,
            Key=key,
        )
        compressed_data = response["Body"].read()
    except ClientError as e:
        raise Exception(f"Failed to download backup: {e}")

    print("Decompressing...")
    sql_data = gzip.decompress(compressed_data)

    # Clean database before restore if requested
    if clean_first:
        clean_database()

    print(f"Restoring to database {config.DB_NAME}...")

    # Build psql command
    env = os.environ.copy()
    env["PGPASSWORD"] = config.DB_PASSWORD

    cmd = [
        "psql",
        "-h",
        config.DB_HOST,
        "-p",
        config.DB_PORT,
        "-U",
        config.DB_USER,
        "-d",
        config.DB_NAME,
    ]

    result = subprocess.run(
        cmd,
        env=env,
        input=sql_data,
        capture_output=True,
    )

    if result.returncode != 0:
        stderr = result.stderr.decode()
        # psql may return warnings that aren't fatal errors
        if "ERROR" in stderr:
            raise Exception(f"psql restore failed: {stderr}")
        else:
            print(f"Warnings: {stderr}")

    print("Restore completed successfully!")


def main() -> int:
    """Main restore routine."""
    # Validate configuration
    if not config.S3_BUCKET_NAME:
        print("Error: S3_BUCKET_NAME is not configured")
        return 1

    s3_client = create_s3_client()

    # Parse arguments
    args = sys.argv[1:]
    clean_first = True

    if "--no-clean" in args:
        clean_first = False
        args.remove("--no-clean")

    if len(args) < 1:
        # List available backups
        backups = list_backups(s3_client)
        if backups:
            print(f"\nTo restore, run: python restore.py <filename>")
            print("Add --no-clean to skip database cleanup before restore")
        else:
            print("No backups found.")
        return 0

    filename = args[0]

    # Confirm restore
    print(f"WARNING: This will restore database from {filename}")
    if clean_first:
        print("Database will be CLEANED (all existing data will be DELETED)!")
    else:
        print("Database will NOT be cleaned (may cause conflicts with existing data)")
    print()

    confirm = input("Type 'yes' to continue: ")

    if confirm.lower() != "yes":
        print("Restore cancelled.")
        return 0

    try:
        restore_backup(s3_client, filename, clean_first=clean_first)
        return 0
    except Exception as e:
        print(f"Restore failed: {e}")
        return 1


if __name__ == "__main__":
    sys.exit(main())