backup.py

#!/usr/bin/env python3
"""
PostgreSQL backup to Beget S3 with password-protected zip compression.
CLI tool for managing automated daily backups with rotation (max 60 copies per DB).
Supports multiple databases.
"""

import argparse
import os
import sys
import subprocess
import tempfile
import logging
from datetime import datetime
from pathlib import Path

from minio import Minio

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

SCRIPT_DIR = Path(__file__).resolve().parent
ENV_FILE = SCRIPT_DIR / ".env"


def load_env(path: Path):
    """Parse a .env file and set environment variables."""
    if not path.exists():
        return
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" not in line:
            continue
        key, _, value = line.partition("=")
        key = key.strip()
        value = value.strip().strip("'\"")
        os.environ.setdefault(key, value)


load_env(ENV_FILE)

DATABASE_URLS = []
for i in range(1, 100):
    url = os.getenv(f"DATABASE_URL_{i}", "")
    if not url:
        break
    DATABASE_URLS.append(url)

S3_ACCESS_KEY_ID = os.getenv("BEGET_S3_ACCESS_KEY_ID", "")
S3_SECRET_ACCESS_KEY = os.getenv("BEGET_S3_SECRET_ACCESS_KEY", "")
S3_BUCKET_NAME = os.getenv("BEGET_S3_BUCKET_NAME", "")
S3_REGION = os.getenv("BEGET_S3_REGION", "ru-1")
S3_ENDPOINT = os.getenv("BEGET_S3_ENDPOINT", "")
BACKUP_PASSWORD = os.getenv("BACKUP_PASSWORD", "renovatio")
MAX_BACKUPS = int(os.getenv("MAX_BACKUPS", "60"))
S3_BACKUP_PREFIX = os.getenv("S3_BACKUP_PREFIX", "backup-db/")

LOG_FORMAT = "%(asctime)s [%(levelname)s] %(message)s"
logging.basicConfig(level=logging.INFO, format=LOG_FORMAT)
log = logging.getLogger("pg-backup")


# ---------------------------------------------------------------------------
# S3 client
# ---------------------------------------------------------------------------

def get_s3_client():
    endpoint = S3_ENDPOINT.replace("https://", "").replace("http://", "")
    return Minio(
        endpoint,
        access_key=S3_ACCESS_KEY_ID,
        secret_key=S3_SECRET_ACCESS_KEY,
        region=S3_REGION,
        secure=True,
    )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def parse_db_url(url: str):
    """Parse postgresql://user:pass@host:port/dbname into components."""
    from urllib.parse import urlparse

    parsed = urlparse(url)
    return {
        "user": parsed.username,
        "password": parsed.password,
        "host": parsed.hostname,
        "port": parsed.port or 5432,
        "dbname": parsed.path.lstrip("/"),
    }


def generate_backup_name(dbname: str) -> str:
    ts = datetime.now().strftime("%Y-%m-%d_%H%M%S")
    return f"{dbname}_{ts}"


def get_db_prefix(dbname: str) -> str:
    """Get the S3 prefix for a specific database."""
    return f"{S3_BACKUP_PREFIX}{dbname}/"


def list_backups(client, dbname: str) -> list:
    """List all backup objects for a specific DB in S3, sorted oldest → newest."""
    prefix = get_db_prefix(dbname)
    objects = list(client.list_objects(S3_BUCKET_NAME, prefix=prefix, recursive=True))
    objects.sort(key=lambda o: o.object_name)
    return objects


def list_all_backups(client) -> dict:
    """List all backup objects grouped by dbname."""
    result = {}
    for url in DATABASE_URLS:
        db = parse_db_url(url)
        result[db["dbname"]] = list_backups(client, db["dbname"])
    return result


def delete_backup(client, key: str):
    client.remove_object(S3_BUCKET_NAME, key)
    log.info(f"Deleted old backup: {key}")


# ---------------------------------------------------------------------------
# Core operations
# ---------------------------------------------------------------------------

def backup_single_db(client, db_url: str):
    """Create a backup for a single DB, zip it with password, upload to S3, rotate."""
    db = parse_db_url(db_url)
    dbname = db["dbname"]
    backup_name = generate_backup_name(dbname)
    s3_key = f"{get_db_prefix(dbname)}{backup_name}.zip"

    log.info(f"Starting backup of database '{dbname}'...")

    with tempfile.TemporaryDirectory() as tmpdir:
        sql_path = os.path.join(tmpdir, f"{backup_name}.sql")
        zip_path = os.path.join(tmpdir, f"{backup_name}.zip")

        # 1. pg_dump
        env = os.environ.copy()
        env["PGPASSWORD"] = db["password"]
        cmd = [
            "pg_dump",
            "-h", db["host"],
            "-p", str(db["port"]),
            "-U", db["user"],
            "-d", dbname,
            "-F", "p",
            "-f", sql_path,
        ]
        log.info(f"Running pg_dump → {sql_path}")
        result = subprocess.run(cmd, env=env, capture_output=True, text=True)
        if result.returncode != 0:
            log.error(f"pg_dump failed:\n{result.stderr}")
            return False

        # 2. zip with password
        log.info(f"Compressing with zip (password-protected) → {zip_path}")
        result = subprocess.run(
            ["zip", "-j", "-P", BACKUP_PASSWORD, zip_path, sql_path],
            capture_output=True, text=True,
        )
        if result.returncode != 0:
            log.error(f"zip failed:\n{result.stderr}")
            return False

        # 3. Upload to S3
        file_size = os.path.getsize(zip_path)
        log.info(f"Uploading to S3: {s3_key} ({file_size / 1024 / 1024:.1f} MB)")
        with open(zip_path, "rb") as f:
            client.put_object(S3_BUCKET_NAME, s3_key, f, file_size)
        log.info(f"Upload complete: {s3_key}")

    # 4. Rotate for this DB
    do_cleanup(client, dbname)

    return True


def do_backup():
    """Create backups for all configured databases, upload to S3, rotate."""
    if not DATABASE_URLS:
        log.error("No DATABASE_URL_* variables found in .env")
        sys.exit(1)

    client = get_s3_client()
    success_count = 0
    for url in DATABASE_URLS:
        db = parse_db_url(url)
        try:
            if backup_single_db(client, url):
                success_count += 1
        except Exception as e:
            log.error(f"Backup of '{db['dbname']}' failed: {e}")

    log.info(f"Backup finished: {success_count}/{len(DATABASE_URLS)} databases succeeded.")
    if success_count == 0:
        sys.exit(1)


def do_list(dbname: str = None):
    """List all backups in S3 for a specific DB or all DBs."""
    client = get_s3_client()
    if dbname:
        backups = list_backups(client, dbname)
        if not backups:
            print(f"No backups found for '{dbname}'.")
            return
        _print_backup_table(dbname, backups)
    else:
        all_backups = list_all_backups(client)
        any_found = False
        for db_name, backups in all_backups.items():
            if backups:
                any_found = True
                _print_backup_table(db_name, backups)
                print()
        if not any_found:
            print("No backups found.")


def _print_backup_table(dbname: str, backups: list):
    print(f"=== {dbname} ===")
    print(f"{'#':<4} {'Filename':<40} {'Size':>12} {'Last Modified'}")
    print("-" * 80)
    for i, obj in enumerate(backups, 1):
        name = os.path.basename(obj.object_name)
        size_mb = obj.size / 1024 / 1024
        modified = obj.last_modified.strftime("%Y-%m-%d %H:%M:%S")
        print(f"{i:<4} {name:<40} {size_mb:>10.1f}MB {modified}")
    print(f"Total: {len(backups)} / {MAX_BACKUPS} backups")


def do_status(dbname: str = None):
    """Show backup count, oldest and newest for a specific DB or all DBs."""
    client = get_s3_client()
    if dbname:
        backups = list_backups(client, dbname)
        if not backups:
            print(f"No backups found for '{dbname}'.")
            return
        _print_status(dbname, backups)
    else:
        all_backups = list_all_backups(client)
        any_found = False
        for db_name, backups in all_backups.items():
            if backups:
                any_found = True
                _print_status(db_name, backups)
                print()
        if not any_found:
            print("No backups found.")


def _print_status(dbname: str, backups: list):
    oldest = backups[0]
    newest = backups[-1]
    total_size = sum(o.size for o in backups) / 1024 / 1024

    print(f"=== {dbname} ===")
    print(f"Backups:  {len(backups)} / {MAX_BACKUPS}")
    print(f"Total:    {total_size:.1f} MB")
    print(f"Oldest:   {os.path.basename(oldest.object_name)}  ({oldest.last_modified.strftime('%Y-%m-%d %H:%M:%S')})")
    print(f"Newest:   {os.path.basename(newest.object_name)}  ({newest.last_modified.strftime('%Y-%m-%d %H:%M:%S')})")


def do_cleanup(client=None, dbname: str = None):
    """Delete oldest backups if count exceeds MAX_BACKUPS for each DB."""
    if client is None:
        client = get_s3_client()

    if dbname:
        _cleanup_db(client, dbname)
    else:
        for url in DATABASE_URLS:
            db = parse_db_url(url)
            _cleanup_db(client, db["dbname"])


def _cleanup_db(client, dbname: str):
    backups = list_backups(client, dbname)
    excess = len(backups) - MAX_BACKUPS
    if excess <= 0:
        log.info(f"Cleanup [{dbname}]: {len(backups)}/{MAX_BACKUPS} backups — no action needed.")
        return

    log.info(f"Cleanup [{dbname}]: {len(backups)} backups, removing {excess} oldest...")
    for obj in backups[:excess]:
        delete_backup(client, obj.object_name)
    log.info(f"Cleanup [{dbname}] done. {len(backups) - excess} backups remaining.")


def do_download(filename: str, dbname: str, output_dir: str = "."):
    """Download a backup from S3 to local directory."""
    client = get_s3_client()
    prefix = get_db_prefix(dbname)
    key = f"{prefix}{filename}" if not filename.startswith(prefix) else filename
    output_path = os.path.join(output_dir, os.path.basename(key))

    log.info(f"Downloading {key} → {output_path}")
    resp = client.get_object(S3_BUCKET_NAME, key)
    with open(output_path, "wb") as f:
        for d in resp.stream(amt=8192):
            f.write(d)
    resp.close()
    log.info(f"Downloaded: {output_path}")
    print(f"Saved to: {output_path}")


def do_restore(filename: str, dbname: str):
    """Download a backup from S3, unzip, and restore to database."""
    db = None
    for url in DATABASE_URLS:
        parsed = parse_db_url(url)
        if parsed["dbname"] == dbname:
            db = parsed
            break
    if db is None:
        log.error(f"Database '{dbname}' not found in configuration.")
        sys.exit(1)

    client = get_s3_client()
    prefix = get_db_prefix(dbname)
    key = f"{prefix}{filename}" if not filename.startswith(prefix) else filename

    with tempfile.TemporaryDirectory() as tmpdir:
        zip_path = os.path.join(tmpdir, os.path.basename(key))
        log.info(f"Downloading {key}...")
        resp = client.get_object(S3_BUCKET_NAME, key)
        with open(zip_path, "wb") as f:
            for d in resp.stream(amt=8192):
                f.write(d)
        resp.close()

        # Unzip with password
        log.info(f"Extracting (password-protected)...")
        result = subprocess.run(
            ["unzip", "-o", "-P", BACKUP_PASSWORD, zip_path, "-d", tmpdir],
            capture_output=True, text=True,
        )
        if result.returncode != 0:
            log.error(f"unzip failed:\n{result.stderr}")
            sys.exit(1)

        # Find the .sql file
        sql_files = list(Path(tmpdir).glob("*.sql"))
        if not sql_files:
            log.error("No .sql file found in archive.")
            sys.exit(1)
        sql_path = str(sql_files[0])

        # Restore via psql
        env = os.environ.copy()
        env["PGPASSWORD"] = db["password"]
        log.info(f"Restoring to database '{db['dbname']}'...")
        cmd = [
            "psql",
            "-h", db["host"],
            "-p", str(db["port"]),
            "-U", db["user"],
            "-d", db["dbname"],
            "-f", sql_path,
        ]
        result = subprocess.run(cmd, env=env, capture_output=True, text=True)
        if result.returncode != 0:
            log.error(f"psql restore failed:\n{result.stderr}")
            sys.exit(1)

    log.info("Restore completed successfully.")


def do_install_cron(hour: int = 3, minute: int = 0):
    """Install a cron job for daily backup."""
    script = str(SCRIPT_DIR / "backup.py")
    log_file = str(SCRIPT_DIR / "backup.log")
    python_bin = str(SCRIPT_DIR / "venv" / "bin" / "python3")
    cron_line = f"{minute} {hour} * * * {python_bin} {script} backup >> {log_file} 2>&1"

    # Read existing crontab
    result = subprocess.run(["crontab", "-l"], capture_output=True, text=True)
    existing = result.stdout.strip() if result.returncode == 0 else ""

    # Remove old entries for this script
    lines = [l for l in existing.splitlines() if "backup.py backup" not in l]
    lines.append(cron_line)
    new_crontab = "\n".join(lines) + "\n"

    subprocess.run(["crontab", "-"], input=new_crontab, text=True, check=True)
    log.info(f"Cron job installed: daily at {hour:02d}:{minute:02d}")
    log.info(f"  {cron_line}")
    print(f"Cron installed: daily at {hour:02d}:{minute:02d}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="PostgreSQL backup to Beget S3 with rotation (multi-DB).",
        prog="backup.py",
    )
    sub = parser.add_subparsers(dest="command", help="Available commands")

    sub.add_parser("backup", help="Create backups for all databases and upload to S3")

    p_list = sub.add_parser("list", help="List all backups in S3")
    p_list.add_argument("--db", default=None, help="Database name (default: show all)")

    p_status = sub.add_parser("status", help="Show backup count, oldest and newest")
    p_status.add_argument("--db", default=None, help="Database name (default: show all)")

    p_cleanup = sub.add_parser("cleanup", help="Delete old backups if count exceeds MAX_BACKUPS")
    p_cleanup.add_argument("--db", default=None, help="Database name (default: cleanup all)")

    p_download = sub.add_parser("download", help="Download a backup from S3")
    p_download.add_argument("filename", help="Backup filename (e.g. blissmeflash_2025-07-27_030000.zip)")
    p_download.add_argument("--db", required=True, help="Database name")
    p_download.add_argument("-o", "--output", default=".", help="Output directory (default: current)")

    p_restore = sub.add_parser("restore", help="Download and restore a backup to the database")
    p_restore.add_argument("filename", help="Backup filename (e.g. blissmeflash_2025-07-27_030000.zip)")
    p_restore.add_argument("--db", required=True, help="Database name")

    p_cron = sub.add_parser("install-cron", help="Install cron job for daily backup")
    p_cron.add_argument("--hour", type=int, default=3, help="Hour (0-23, default: 3)")
    p_cron.add_argument("--minute", type=int, default=0, help="Minute (0-59, default: 0)")

    args = parser.parse_args()

    if args.command is None:
        parser.print_help()
        sys.exit(0)

    if args.command == "backup":
        do_backup()
    elif args.command == "list":
        do_list(args.db)
    elif args.command == "status":
        do_status(args.db)
    elif args.command == "cleanup":
        do_cleanup(dbname=args.db)
    elif args.command == "download":
        do_download(args.filename, args.db, args.output)
    elif args.command == "restore":
        do_restore(args.filename, args.db)
    elif args.command == "install-cron":
        do_install_cron(args.hour, args.minute)


if __name__ == "__main__":
    main()