159 lines
4.0 KiB
Python
159 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Restore PostgreSQL database from S3 backup.
|
|
|
|
Usage:
|
|
python restore.py - List available backups
|
|
python restore.py <filename> - Restore from specific backup
|
|
"""
|
|
import gzip
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
|
|
import boto3
|
|
from botocore.config import Config as BotoConfig
|
|
from botocore.exceptions import ClientError
|
|
|
|
from config import config
|
|
|
|
|
|
def create_s3_client():
|
|
"""Initialize S3 client."""
|
|
return boto3.client(
|
|
"s3",
|
|
endpoint_url=config.S3_ENDPOINT_URL,
|
|
aws_access_key_id=config.S3_ACCESS_KEY_ID,
|
|
aws_secret_access_key=config.S3_SECRET_ACCESS_KEY,
|
|
region_name=config.S3_REGION or "us-east-1",
|
|
config=BotoConfig(signature_version="s3v4"),
|
|
)
|
|
|
|
|
|
def list_backups(s3_client) -> list[tuple[str, float, str]]:
|
|
"""List all available backups."""
|
|
print("Available backups:\n")
|
|
|
|
try:
|
|
paginator = s3_client.get_paginator("list_objects_v2")
|
|
pages = paginator.paginate(
|
|
Bucket=config.S3_BUCKET_NAME,
|
|
Prefix=config.S3_BACKUP_PREFIX,
|
|
)
|
|
|
|
backups = []
|
|
for page in pages:
|
|
for obj in page.get("Contents", []):
|
|
filename = obj["Key"].replace(config.S3_BACKUP_PREFIX, "")
|
|
size_mb = obj["Size"] / (1024 * 1024)
|
|
modified = obj["LastModified"].strftime("%Y-%m-%d %H:%M:%S")
|
|
backups.append((filename, size_mb, modified))
|
|
|
|
# Sort by date descending (newest first)
|
|
backups.sort(key=lambda x: x[2], reverse=True)
|
|
|
|
for filename, size_mb, modified in backups:
|
|
print(f" {filename} ({size_mb:.2f} MB) - {modified}")
|
|
|
|
return backups
|
|
|
|
except ClientError as e:
|
|
print(f"Error listing backups: {e}")
|
|
return []
|
|
|
|
|
|
def restore_backup(s3_client, filename: str) -> None:
|
|
"""Download and restore backup."""
|
|
key = f"{config.S3_BACKUP_PREFIX}{filename}"
|
|
|
|
print(f"Downloading {filename} from S3...")
|
|
try:
|
|
response = s3_client.get_object(
|
|
Bucket=config.S3_BUCKET_NAME,
|
|
Key=key,
|
|
)
|
|
compressed_data = response["Body"].read()
|
|
except ClientError as e:
|
|
raise Exception(f"Failed to download backup: {e}")
|
|
|
|
print("Decompressing...")
|
|
sql_data = gzip.decompress(compressed_data)
|
|
|
|
print(f"Restoring to database {config.DB_NAME}...")
|
|
|
|
# Build psql command
|
|
env = os.environ.copy()
|
|
env["PGPASSWORD"] = config.DB_PASSWORD
|
|
|
|
cmd = [
|
|
"psql",
|
|
"-h",
|
|
config.DB_HOST,
|
|
"-p",
|
|
config.DB_PORT,
|
|
"-U",
|
|
config.DB_USER,
|
|
"-d",
|
|
config.DB_NAME,
|
|
]
|
|
|
|
result = subprocess.run(
|
|
cmd,
|
|
env=env,
|
|
input=sql_data,
|
|
capture_output=True,
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
stderr = result.stderr.decode()
|
|
# psql may return warnings that aren't fatal errors
|
|
if "ERROR" in stderr:
|
|
raise Exception(f"psql restore failed: {stderr}")
|
|
else:
|
|
print(f"Warnings: {stderr}")
|
|
|
|
print("Restore completed successfully!")
|
|
|
|
|
|
def main() -> int:
|
|
"""Main restore routine."""
|
|
# Validate configuration
|
|
if not config.S3_BUCKET_NAME:
|
|
print("Error: S3_BUCKET_NAME is not configured")
|
|
return 1
|
|
|
|
s3_client = create_s3_client()
|
|
|
|
if len(sys.argv) < 2:
|
|
# List available backups
|
|
backups = list_backups(s3_client)
|
|
if backups:
|
|
print(f"\nTo restore, run: python restore.py <filename>")
|
|
else:
|
|
print("No backups found.")
|
|
return 0
|
|
|
|
filename = sys.argv[1]
|
|
|
|
# Confirm restore
|
|
print(f"WARNING: This will restore database from {filename}")
|
|
print("This may overwrite existing data!")
|
|
print()
|
|
|
|
confirm = input("Type 'yes' to continue: ")
|
|
|
|
if confirm.lower() != "yes":
|
|
print("Restore cancelled.")
|
|
return 0
|
|
|
|
try:
|
|
restore_backup(s3_client, filename)
|
|
return 0
|
|
except Exception as e:
|
|
print(f"Restore failed: {e}")
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|