| import argparse |
| import subprocess |
| import os |
| import sys |
| from datetime import datetime |
|
|
| def backup_database(db_url: str, output_dir: str): |
| os.makedirs(output_dir, exist_ok=True) |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| backup_file = os.path.join(output_dir, f"db_backup_{timestamp}.sql") |
| |
| print(f"[*] Starting database dump to {backup_file}...") |
| cmd = f"pg_dump {db_url} -F p -f {backup_file}" |
| result = subprocess.run(cmd, shell=True, capture_output=True, text=True) |
| |
| if result.returncode == 0: |
| print(f"[✓] Backup completed successfully: {backup_file}") |
| else: |
| print(f"[✗] Backup failed: {result.stderr}") |
|
|
| def restore_database(db_url: str, backup_file: str): |
| if not os.path.exists(backup_file): |
| print(f"[✗] Backup file not found: {backup_file}") |
| return |
|
|
| print(f"[*] Restoring database from {backup_file}...") |
| cmd = f"psql {db_url} -f {backup_file}" |
| result = subprocess.run(cmd, shell=True, capture_output=True, text=True) |
| |
| if result.returncode == 0: |
| print("[✓] Restoration completed successfully.") |
| else: |
| print(f"[✗] Restoration failed: {result.stderr}") |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Database Backup & Restore Utility") |
| parser.add_argument("action", choices=["backup", "restore"], help="Action to execute") |
| parser.add_argument("--db-url", required=True, help="PostgreSQL connection string") |
| parser.add_argument("--dir", default="./backups", help="Backup directory") |
| parser.add_argument("--file", help="Backup file path (required for restore)") |
|
|
| args = parser.parse_args() |
|
|
| if args.action == "backup": |
| backup_database(args.db_url, args.dir) |
| elif args.action == "restore": |
| if not args.file: |
| print("[✗] Error: --file argument is required for restore.") |
| sys.exit(1) |
| restore_database(args.db_url, args.file) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
