team-tryouts/app/supporting_scripts/backup.py

"""PostgreSQL database backup and restore helpers for the admin panel.

Unlike the original SQLite-based script, this module works against the
PostgreSQL database configured via DATABASE_URL. It shells out to the
standard ``pg_dump`` / ``pg_restore`` / ``psql`` command-line utilities.

Configuration via environment variables:
    BACKUP_DIR: Directory to store backups (default: ./backups)
    BACKUP_RETENTION_DAYS: Number of days to keep backups (default: 30)
    DATABASE_URL: PostgreSQL connection string (required for connectivity)
"""

import os
import re
import subprocess
from datetime import datetime, timedelta
from urllib.parse import urlparse

BACKUP_DIR = os.getenv('BACKUP_DIR', os.path.join(os.getcwd(), 'backups'))
BACKUP_RETENTION_DAYS = int(os.getenv('BACKUP_RETENTION_DAYS', 30))


def _database_url():
    """Return the DATABASE_URL, normalizing to the psycopg2 SQLAlchemy form."""
    url = os.getenv('DATABASE_URL')
    if not url:
        raise RuntimeError('DATABASE_URL environment variable must be set')
    # SQLAlchemy may prefix driver; strip it for command-line tools.
    return re.sub(r'^postgresql\+[^:]+://', 'postgresql://', url)


def _pg_env():
    """Return the dict of PG* variables required by pg_dump / psql."""
    parsed = urlparse(_database_url())
    env = os.environ.copy()
    # PGPASSWORD avoids interactive password prompts.
    if parsed.password is not None:
        env['PGPASSWORD'] = parsed.password
    env.pop('PGHOST', None)
    env.pop('PGPORT', None)
    env.pop('PGUSER', None)
    env.pop('PGDATABASE', None)
    return env


def _pg_args():
    """Return [host, port, user, dbname] positional args for pg tools."""
    parsed = urlparse(_database_url())
    args = []
    if parsed.hostname:
        args += ['--host', parsed.hostname]
    if parsed.port:
        args += ['--port', str(parsed.port)]
    if parsed.username:
        args += ['--username', parsed.username]
    if parsed.path and parsed.path.lstrip('/'):
        args += ['--dbname', parsed.path.lstrip('/')]
    return args


def create_backup_dir():
    """Create the backup directory if it doesn't exist."""
    os.makedirs(BACKUP_DIR, exist_ok=True)


def create_backup(backup_type='manual', notes=None):
    """Create a PostgreSQL dump backup.

    Returns:
        dict: {'filename', 'file_path', 'size_bytes'} or None on failure.
    """
    create_backup_dir()
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    filename = f'db_backup_{backup_type}_{timestamp}.sql'
    file_path = os.path.join(BACKUP_DIR, filename)

    cmd = ['pg_dump'] + _pg_args() + ['--format', 'plain', '--no-owner', '--file', file_path]
    try:
        result = subprocess.run(cmd, env=_pg_env(), capture_output=True, text=True)
        if result.returncode != 0:
            _cleanup_failed_file(file_path)
            raise RuntimeError(result.stderr.strip() or 'pg_dump failed')
    except FileNotFoundError:
        _cleanup_failed_file(file_path)
        raise RuntimeError('pg_dump executable not found on PATH')
    except Exception:
        _cleanup_failed_file(file_path)
        raise

    if not os.path.exists(file_path):
        raise RuntimeError('Backup file was not created')

    size = os.path.getsize(file_path)
    return {'filename': filename, 'file_path': file_path, 'size_bytes': size}


def _cleanup_failed_file(file_path):
    """Remove a partially-written dump file, ignoring errors."""
    if file_path and os.path.exists(file_path):
        try:
            os.remove(file_path)
        except OSError:
            pass


def restore_backup(file_path):
    """Restore a PostgreSQL dump file using psql.

    Args:
        file_path: Path to a plain-text .sql dump produced by pg_dump.

    Raises:
        RuntimeError: If psql fails or is unavailable.
    """
    if not file_path or not os.path.exists(file_path):
        raise RuntimeError('Backup file does not exist')

    # The dump was produced with --no-owner + plain format, so we can
    # pipe it into psql. Use --set ON_ERROR_STOP=1 to fail fast.
    cmd = ['psql'] + _pg_args() + ['--set', 'ON_ERROR_STOP=1', '--file', file_path]
    try:
        result = subprocess.run(cmd, env=_pg_env(), capture_output=True, text=True)
        if result.returncode != 0:
            raise RuntimeError(result.stderr.strip() or 'psql restore failed')
    except FileNotFoundError:
        raise RuntimeError('psql executable not found on PATH')


def cleanup_old_backups():
    """Remove backup files older than BACKUP_RETENTION_DAYS."""
    if not os.path.exists(BACKUP_DIR):
        return 0

    cutoff = datetime.now() - timedelta(days=BACKUP_RETENTION_DAYS)
    removed_count = 0
    for filename in os.listdir(BACKUP_DIR):
        file_path = os.path.join(BACKUP_DIR, filename)
        if not os.path.isfile(file_path):
            continue
        file_time = datetime.fromtimestamp(os.path.getmtime(file_path))
        if file_time < cutoff:
            try:
                os.remove(file_path)
                removed_count += 1
            except OSError:
                pass
    return removed_count


def main():
    """Run a manual backup from the CLI (for cron/Task Scheduler)."""
    print('=== Team Tryouts PostgreSQL Backup ===')
    print(f'Started at: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}')
    print(f'Backup directory: {BACKUP_DIR}')
    print(f'Retention period: {BACKUP_RETENTION_DAYS} days')
    print()

    try:
        result = create_backup(backup_type='scheduled')
    except RuntimeError as e:
        print(f'[ERROR] {e}')
        return 1

    print(f'[OK] Database backed up to: {result["file_path"]}')

    removed = cleanup_old_backups()
    if removed:
        print(f'[CLEANUP] Removed {removed} old backup(s).')

    print()
    print('=== Backup completed successfully ===')
    return 0


if __name__ == '__main__':
    exit(main())