Initial commit
This commit is contained in:
Zhongwei Li
502
skills/databases/scripts/db_backup.py
Normal file
502
skills/databases/scripts/db_backup.py
Normal file
@@ -0,0 +1,502 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Database backup and restore tool for MongoDB and PostgreSQL.
|
||||
Supports compression, scheduling, and verification.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gzip
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class BackupInfo:
|
||||
"""Backup metadata."""
|
||||
|
||||
filename: str
|
||||
database_type: str
|
||||
database_name: str
|
||||
timestamp: datetime
|
||||
size_bytes: int
|
||||
compressed: bool
|
||||
verified: bool = False
|
||||
|
||||
|
||||
class BackupManager:
|
||||
"""Manages database backups for MongoDB and PostgreSQL."""
|
||||
|
||||
def __init__(self, db_type: str, backup_dir: str = "./backups"):
|
||||
"""
|
||||
Initialize backup manager.
|
||||
|
||||
Args:
|
||||
db_type: Database type ('mongodb' or 'postgres')
|
||||
backup_dir: Directory to store backups
|
||||
"""
|
||||
self.db_type = db_type.lower()
|
||||
self.backup_dir = Path(backup_dir)
|
||||
self.backup_dir.mkdir(exist_ok=True)
|
||||
|
||||
def create_backup(
|
||||
self,
|
||||
uri: str,
|
||||
database: Optional[str] = None,
|
||||
compress: bool = True,
|
||||
verify: bool = True
|
||||
) -> Optional[BackupInfo]:
|
||||
"""
|
||||
Create database backup.
|
||||
|
||||
Args:
|
||||
uri: Database connection string
|
||||
database: Database name (optional for MongoDB)
|
||||
compress: Compress backup file
|
||||
verify: Verify backup after creation
|
||||
|
||||
Returns:
|
||||
BackupInfo if successful, None otherwise
|
||||
"""
|
||||
timestamp = datetime.now()
|
||||
date_str = timestamp.strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
if self.db_type == "mongodb":
|
||||
return self._backup_mongodb(uri, database, date_str, compress, verify)
|
||||
elif self.db_type == "postgres":
|
||||
return self._backup_postgres(uri, database, date_str, compress, verify)
|
||||
else:
|
||||
print(f"Error: Unsupported database type: {self.db_type}")
|
||||
return None
|
||||
|
||||
def _backup_mongodb(
|
||||
self,
|
||||
uri: str,
|
||||
database: Optional[str],
|
||||
date_str: str,
|
||||
compress: bool,
|
||||
verify: bool
|
||||
) -> Optional[BackupInfo]:
|
||||
"""Create MongoDB backup using mongodump."""
|
||||
db_name = database or "all"
|
||||
filename = f"mongodb_{db_name}_{date_str}"
|
||||
backup_path = self.backup_dir / filename
|
||||
|
||||
try:
|
||||
cmd = ["mongodump", "--uri", uri, "--out", str(backup_path)]
|
||||
|
||||
if database:
|
||||
cmd.extend(["--db", database])
|
||||
|
||||
print(f"Creating MongoDB backup: {filename}")
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f"Error: {result.stderr}")
|
||||
return None
|
||||
|
||||
# Compress if requested
|
||||
if compress:
|
||||
archive_path = backup_path.with_suffix(".tar.gz")
|
||||
print(f"Compressing backup...")
|
||||
shutil.make_archive(str(backup_path), "gztar", backup_path)
|
||||
shutil.rmtree(backup_path)
|
||||
backup_path = archive_path
|
||||
filename = archive_path.name
|
||||
|
||||
size_bytes = self._get_size(backup_path)
|
||||
|
||||
backup_info = BackupInfo(
|
||||
filename=filename,
|
||||
database_type="mongodb",
|
||||
database_name=db_name,
|
||||
timestamp=datetime.now(),
|
||||
size_bytes=size_bytes,
|
||||
compressed=compress
|
||||
)
|
||||
|
||||
if verify:
|
||||
backup_info.verified = self._verify_backup(backup_info)
|
||||
|
||||
self._save_metadata(backup_info)
|
||||
print(f"✓ Backup created: {filename} ({self._format_size(size_bytes)})")
|
||||
|
||||
return backup_info
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error creating MongoDB backup: {e}")
|
||||
return None
|
||||
|
||||
def _backup_postgres(
|
||||
self,
|
||||
uri: str,
|
||||
database: str,
|
||||
date_str: str,
|
||||
compress: bool,
|
||||
verify: bool
|
||||
) -> Optional[BackupInfo]:
|
||||
"""Create PostgreSQL backup using pg_dump."""
|
||||
if not database:
|
||||
print("Error: Database name required for PostgreSQL backup")
|
||||
return None
|
||||
|
||||
ext = ".sql.gz" if compress else ".sql"
|
||||
filename = f"postgres_{database}_{date_str}{ext}"
|
||||
backup_path = self.backup_dir / filename
|
||||
|
||||
try:
|
||||
cmd = ["pg_dump", uri]
|
||||
|
||||
if compress:
|
||||
# Use pg_dump with gzip
|
||||
with open(backup_path, "wb") as f:
|
||||
dump_proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
|
||||
gzip_proc = subprocess.Popen(
|
||||
["gzip"],
|
||||
stdin=dump_proc.stdout,
|
||||
stdout=f
|
||||
)
|
||||
dump_proc.stdout.close()
|
||||
gzip_proc.communicate()
|
||||
|
||||
if dump_proc.returncode != 0:
|
||||
print("Error: pg_dump failed")
|
||||
return None
|
||||
else:
|
||||
with open(backup_path, "w") as f:
|
||||
result = subprocess.run(cmd, stdout=f, stderr=subprocess.PIPE, text=True)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f"Error: {result.stderr}")
|
||||
return None
|
||||
|
||||
size_bytes = backup_path.stat().st_size
|
||||
|
||||
backup_info = BackupInfo(
|
||||
filename=filename,
|
||||
database_type="postgres",
|
||||
database_name=database,
|
||||
timestamp=datetime.now(),
|
||||
size_bytes=size_bytes,
|
||||
compressed=compress
|
||||
)
|
||||
|
||||
if verify:
|
||||
backup_info.verified = self._verify_backup(backup_info)
|
||||
|
||||
self._save_metadata(backup_info)
|
||||
print(f"✓ Backup created: {filename} ({self._format_size(size_bytes)})")
|
||||
|
||||
return backup_info
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error creating PostgreSQL backup: {e}")
|
||||
return None
|
||||
|
||||
def restore_backup(self, filename: str, uri: str, dry_run: bool = False) -> bool:
|
||||
"""
|
||||
Restore database from backup.
|
||||
|
||||
Args:
|
||||
filename: Backup filename
|
||||
uri: Database connection string
|
||||
dry_run: If True, only show what would be done
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
backup_path = self.backup_dir / filename
|
||||
|
||||
if not backup_path.exists():
|
||||
print(f"Error: Backup not found: {filename}")
|
||||
return False
|
||||
|
||||
# Load metadata
|
||||
metadata_path = backup_path.with_suffix(".json")
|
||||
if metadata_path.exists():
|
||||
with open(metadata_path) as f:
|
||||
metadata = json.load(f)
|
||||
print(f"Restoring backup from {metadata['timestamp']}")
|
||||
print(f"Database: {metadata['database_name']}")
|
||||
|
||||
if dry_run:
|
||||
print(f"Would restore from: {backup_path}")
|
||||
return True
|
||||
|
||||
print(f"Restoring backup: {filename}")
|
||||
|
||||
try:
|
||||
if self.db_type == "mongodb":
|
||||
return self._restore_mongodb(backup_path, uri)
|
||||
elif self.db_type == "postgres":
|
||||
return self._restore_postgres(backup_path, uri)
|
||||
else:
|
||||
print(f"Error: Unsupported database type: {self.db_type}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error restoring backup: {e}")
|
||||
return False
|
||||
|
||||
def _restore_mongodb(self, backup_path: Path, uri: str) -> bool:
|
||||
"""Restore MongoDB backup using mongorestore."""
|
||||
try:
|
||||
# Extract if compressed
|
||||
restore_path = backup_path
|
||||
if backup_path.suffix == ".gz":
|
||||
print("Extracting backup...")
|
||||
extract_path = backup_path.with_suffix("")
|
||||
shutil.unpack_archive(backup_path, extract_path)
|
||||
restore_path = extract_path
|
||||
|
||||
cmd = ["mongorestore", "--uri", uri, str(restore_path)]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
|
||||
# Cleanup extracted files
|
||||
if restore_path != backup_path and restore_path.is_dir():
|
||||
shutil.rmtree(restore_path)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f"Error: {result.stderr}")
|
||||
return False
|
||||
|
||||
print("✓ Restore completed")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error restoring MongoDB: {e}")
|
||||
return False
|
||||
|
||||
def _restore_postgres(self, backup_path: Path, uri: str) -> bool:
|
||||
"""Restore PostgreSQL backup using psql."""
|
||||
try:
|
||||
if backup_path.suffix == ".gz":
|
||||
# Decompress and restore
|
||||
with gzip.open(backup_path, "rb") as f:
|
||||
cmd = ["psql", uri]
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
stdin=f,
|
||||
capture_output=True,
|
||||
text=False
|
||||
)
|
||||
else:
|
||||
with open(backup_path) as f:
|
||||
cmd = ["psql", uri]
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
stdin=f,
|
||||
capture_output=True,
|
||||
text=True
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f"Error: {result.stderr}")
|
||||
return False
|
||||
|
||||
print("✓ Restore completed")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error restoring PostgreSQL: {e}")
|
||||
return False
|
||||
|
||||
def list_backups(self) -> List[BackupInfo]:
|
||||
"""
|
||||
List all backups.
|
||||
|
||||
Returns:
|
||||
List of BackupInfo objects
|
||||
"""
|
||||
backups = []
|
||||
|
||||
for metadata_file in sorted(self.backup_dir.glob("*.json")):
|
||||
try:
|
||||
with open(metadata_file) as f:
|
||||
data = json.load(f)
|
||||
|
||||
backup_info = BackupInfo(
|
||||
filename=data["filename"],
|
||||
database_type=data["database_type"],
|
||||
database_name=data["database_name"],
|
||||
timestamp=datetime.fromisoformat(data["timestamp"]),
|
||||
size_bytes=data["size_bytes"],
|
||||
compressed=data["compressed"],
|
||||
verified=data.get("verified", False)
|
||||
)
|
||||
backups.append(backup_info)
|
||||
except Exception as e:
|
||||
print(f"Error reading metadata {metadata_file}: {e}")
|
||||
|
||||
return backups
|
||||
|
||||
def cleanup_old_backups(self, retention_days: int, dry_run: bool = False) -> int:
|
||||
"""
|
||||
Remove backups older than retention period.
|
||||
|
||||
Args:
|
||||
retention_days: Number of days to retain backups
|
||||
dry_run: If True, only show what would be deleted
|
||||
|
||||
Returns:
|
||||
Number of backups removed
|
||||
"""
|
||||
cutoff = datetime.now().timestamp() - (retention_days * 24 * 3600)
|
||||
removed = 0
|
||||
|
||||
for backup_file in self.backup_dir.glob("*"):
|
||||
if backup_file.suffix == ".json":
|
||||
continue
|
||||
|
||||
if backup_file.stat().st_mtime < cutoff:
|
||||
if dry_run:
|
||||
print(f"Would remove: {backup_file.name}")
|
||||
else:
|
||||
print(f"Removing: {backup_file.name}")
|
||||
backup_file.unlink()
|
||||
# Remove metadata
|
||||
metadata_file = backup_file.with_suffix(".json")
|
||||
if metadata_file.exists():
|
||||
metadata_file.unlink()
|
||||
removed += 1
|
||||
|
||||
return removed
|
||||
|
||||
def _verify_backup(self, backup_info: BackupInfo) -> bool:
|
||||
"""
|
||||
Verify backup integrity.
|
||||
|
||||
Args:
|
||||
backup_info: Backup information
|
||||
|
||||
Returns:
|
||||
True if backup is valid, False otherwise
|
||||
"""
|
||||
backup_path = self.backup_dir / backup_info.filename
|
||||
|
||||
if not backup_path.exists():
|
||||
return False
|
||||
|
||||
# Basic verification: file exists and has size > 0
|
||||
if backup_path.stat().st_size == 0:
|
||||
return False
|
||||
|
||||
# Could add more verification here (checksums, test restore, etc.)
|
||||
return True
|
||||
|
||||
def _get_size(self, path: Path) -> int:
|
||||
"""Get total size of file or directory."""
|
||||
if path.is_file():
|
||||
return path.stat().st_size
|
||||
elif path.is_dir():
|
||||
total = 0
|
||||
for item in path.rglob("*"):
|
||||
if item.is_file():
|
||||
total += item.stat().st_size
|
||||
return total
|
||||
return 0
|
||||
|
||||
def _format_size(self, size_bytes: int) -> str:
|
||||
"""Format size in human-readable format."""
|
||||
for unit in ["B", "KB", "MB", "GB", "TB"]:
|
||||
if size_bytes < 1024:
|
||||
return f"{size_bytes:.2f} {unit}"
|
||||
size_bytes /= 1024
|
||||
return f"{size_bytes:.2f} PB"
|
||||
|
||||
def _save_metadata(self, backup_info: BackupInfo):
|
||||
"""Save backup metadata to JSON file."""
|
||||
metadata_path = self.backup_dir / f"{backup_info.filename}.json"
|
||||
|
||||
metadata = {
|
||||
"filename": backup_info.filename,
|
||||
"database_type": backup_info.database_type,
|
||||
"database_name": backup_info.database_name,
|
||||
"timestamp": backup_info.timestamp.isoformat(),
|
||||
"size_bytes": backup_info.size_bytes,
|
||||
"compressed": backup_info.compressed,
|
||||
"verified": backup_info.verified
|
||||
}
|
||||
|
||||
with open(metadata_path, "w") as f:
|
||||
json.dump(metadata, f, indent=2)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point."""
|
||||
parser = argparse.ArgumentParser(description="Database backup tool")
|
||||
parser.add_argument("--db", required=True, choices=["mongodb", "postgres"],
|
||||
help="Database type")
|
||||
parser.add_argument("--backup-dir", default="./backups",
|
||||
help="Backup directory")
|
||||
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
# Backup command
|
||||
backup_parser = subparsers.add_parser("backup", help="Create backup")
|
||||
backup_parser.add_argument("--uri", required=True, help="Database connection string")
|
||||
backup_parser.add_argument("--database", help="Database name")
|
||||
backup_parser.add_argument("--no-compress", action="store_true",
|
||||
help="Disable compression")
|
||||
backup_parser.add_argument("--no-verify", action="store_true",
|
||||
help="Skip verification")
|
||||
|
||||
# Restore command
|
||||
restore_parser = subparsers.add_parser("restore", help="Restore backup")
|
||||
restore_parser.add_argument("filename", help="Backup filename")
|
||||
restore_parser.add_argument("--uri", required=True, help="Database connection string")
|
||||
restore_parser.add_argument("--dry-run", action="store_true",
|
||||
help="Show what would be done")
|
||||
|
||||
# List command
|
||||
subparsers.add_parser("list", help="List backups")
|
||||
|
||||
# Cleanup command
|
||||
cleanup_parser = subparsers.add_parser("cleanup", help="Remove old backups")
|
||||
cleanup_parser.add_argument("--retention-days", type=int, default=7,
|
||||
help="Days to retain backups (default: 7)")
|
||||
cleanup_parser.add_argument("--dry-run", action="store_true",
|
||||
help="Show what would be removed")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
manager = BackupManager(args.db, args.backup_dir)
|
||||
|
||||
if args.command == "backup":
|
||||
backup_info = manager.create_backup(
|
||||
args.uri,
|
||||
args.database,
|
||||
compress=not args.no_compress,
|
||||
verify=not args.no_verify
|
||||
)
|
||||
sys.exit(0 if backup_info else 1)
|
||||
|
||||
elif args.command == "restore":
|
||||
success = manager.restore_backup(args.filename, args.uri, args.dry_run)
|
||||
sys.exit(0 if success else 1)
|
||||
|
||||
elif args.command == "list":
|
||||
backups = manager.list_backups()
|
||||
print(f"Total backups: {len(backups)}\n")
|
||||
for backup in backups:
|
||||
verified_str = "✓" if backup.verified else "?"
|
||||
print(f"[{verified_str}] {backup.filename}")
|
||||
print(f" Database: {backup.database_name}")
|
||||
print(f" Created: {backup.timestamp}")
|
||||
print(f" Size: {manager._format_size(backup.size_bytes)}")
|
||||
print()
|
||||
|
||||
elif args.command == "cleanup":
|
||||
removed = manager.cleanup_old_backups(args.retention_days, args.dry_run)
|
||||
print(f"Removed {removed} backup(s)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user