mysql backup tool

script/database/mysql_dump_table.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+"""
+Back up a single production database table to a local SQL file using mysqldump.
+
+Reads connection info from secrets.toml (same format as the other database scripts).
+
+Usage:
+  python mysql_dump_table.py --table=users
+  python mysql_dump_table.py --table=media_files --directory=/tmp/backups
+
+Requirements:
+  pip install pymysql
+  mysqldump must be on PATH
+"""
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import tomllib
+from datetime import datetime
+from pathlib import Path
+from urllib.parse import urlparse
+
+try:
+    import pymysql
+except ImportError:
+    print("pymysql is required: pip install pymysql")
+    sys.exit(1)
+
+
+SECRETS_PATH = Path(__file__).parent / "secrets.toml"
+
+DEFAULT_BACKUP_DIR = Path.home() / "database" / "backups"
+
+ALLOWED_TABLES = [
+    "audit_logs",
+    "batch_generations",
+    "favorites",
+    "featured_items",
+    "generic_inference_jobs",
+    "google_sign_in_accounts",
+    "media_files",
+    "model_weights",
+    "prompt_context_items",
+    "prompts",
+    "tag_uses",
+    "tags",
+    "user_bookmarks",
+    "user_ratings",
+    "user_roles",
+    "user_sessions",
+    "user_stripe_customer_links",
+    "user_subscriptions",
+    "users",
+    "wallet_ledger_entries",
+    "wallets",
+]
+
+
+def load_config():
+    if not SECRETS_PATH.exists():
+        print(f"Error: secrets.toml not found at {SECRETS_PATH}")
+        print(f"Copy secrets.toml.example to secrets.toml and fill in the connection URL.")
+        sys.exit(1)
+
+    with open(SECRETS_PATH, "rb") as f:
+        config = tomllib.load(f)
+    url = config["database"]["url"]
+    parsed = urlparse(url)
+    return {
+        "host": parsed.hostname,
+        "port": parsed.port or 3306,
+        "user": parsed.username,
+        "password": parsed.password,
+        "database": parsed.path.lstrip("/"),
+    }
+
+
+def connect(config):
+    return pymysql.connect(
+        host=config["host"],
+        port=config["port"],
+        user=config["user"],
+        password=config["password"],
+        database=config["database"],
+        charset="utf8mb4",
+        cursorclass=pymysql.cursors.DictCursor,
+    )
+
+
+def table_exists(conn, table_name):
+    with conn.cursor() as cursor:
+        cursor.execute("SHOW TABLES LIKE %s", (table_name,))
+        return cursor.fetchone() is not None
+
+
+def count_rows(conn, table_name):
+    with conn.cursor() as cursor:
+        cursor.execute(f"SELECT COUNT(*) AS cnt FROM `{table_name}`")
+        result = cursor.fetchone()
+        return result["cnt"]
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Back up a database table using mysqldump.")
+    parser.add_argument("--table", required=True, help="Name of the table to back up.")
+    parser.add_argument("--directory", default=None, help=f"Directory to save the dump. Default: {DEFAULT_BACKUP_DIR}")
+    args = parser.parse_args()
+
+    table_name = args.table
+    backup_dir = Path(args.directory) if args.directory else DEFAULT_BACKUP_DIR
+
+    # Validate table is in allowed list
+    if table_name not in ALLOWED_TABLES:
+        print(f"Error: '{table_name}' is not in the allowed tables list.")
+        print(f"Allowed tables: {', '.join(sorted(ALLOWED_TABLES))}")
+        sys.exit(1)
+
+    # Check mysqldump is available
+    if shutil.which("mysqldump") is None:
+        print("Error: mysqldump not found on PATH.")
+        sys.exit(1)
+
+    config = load_config()
+
+    # Verify table exists and count rows
+    conn = connect(config)
+    try:
+        if not table_exists(conn, table_name):
+            print(f"Error: table '{table_name}' does not exist in database '{config['database']}'.")
+            sys.exit(1)
+
+        row_count = count_rows(conn, table_name)
+    finally:
+        conn.close()
+
+    # Build output path
+    backup_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+    output_file = backup_dir / f"{table_name}_{timestamp}.sql"
+
+    # Confirm with user
+    print(f"Table:     {table_name}")
+    print(f"Database:  {config['database']} @ {config['host']}:{config['port']}")
+    print(f"Rows:      {row_count:,}")
+    print(f"Output:    {output_file}")
+    print()
+
+    confirm = input("Proceed with backup? [y/N] ").strip().lower()
+    if confirm not in ("y", "yes"):
+        print("Aborted.")
+        sys.exit(0)
+
+    # Run mysqldump
+    cmd = [
+        "mysqldump",
+        f"--host={config['host']}",
+        f"--port={config['port']}",
+        f"--user={config['user']}",
+        f"--password={config['password']}",
+        "--single-transaction",
+        "--routines",
+        "--triggers",
+        "--set-gtid-purged=OFF",
+        config["database"],
+        table_name,
+    ]
+
+    print(f"\nRunning mysqldump for '{table_name}'...")
+
+    with open(output_file, "w") as f:
+        result = subprocess.run(cmd, stdout=f, stderr=subprocess.PIPE, text=True)
+
+    if result.returncode != 0:
+        print(f"Error: mysqldump failed (exit code {result.returncode})")
+        if result.stderr:
+            # Filter out the password warning which is expected
+            for line in result.stderr.strip().splitlines():
+                if "Using a password on the command line" not in line:
+                    print(f"  {line}")
+        sys.exit(1)
+
+    file_size = output_file.stat().st_size
+    if file_size < 100:
+        print(f"Warning: output file is suspiciously small ({file_size} bytes).")
+
+    print(f"Done. Backup saved to: {output_file} ({file_size:,} bytes)")
+
+
+if __name__ == "__main__":
+    main()