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 specific backup
"""
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 restore_backup(s3_client, filename: str) -> 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)

    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()

    if len(sys.argv) < 2:
        # List available backups
        backups = list_backups(s3_client)
        if backups:
            print(f"\nTo restore, run: python restore.py <filename>")
        else:
            print("No backups found.")
        return 0

    filename = sys.argv[1]

    # Confirm restore
    print(f"WARNING: This will restore database from {filename}")
    print("This may overwrite existing data!")
    print()

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

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

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


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