Agent Zero
272 lines
7.5 KiB
Python
272 lines
7.5 KiB
Python
"""Service for database backup and restore operations."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.backup import Backup
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BACKUP_DIR = Path("/data/backups")
|
|
|
|
|
|
def _ensure_backup_dir() -> None:
|
|
"""Ensure the backup directory exists."""
|
|
BACKUP_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def _get_database_url() -> str:
|
|
"""Get DATABASE_URL from environment."""
|
|
url = os.environ.get("DATABASE_URL")
|
|
if not url:
|
|
raise RuntimeError("DATABASE_URL environment variable is not set")
|
|
return url
|
|
|
|
|
|
def _parse_pg_url(url: str) -> dict[str, str]:
|
|
"""Parse a PostgreSQL connection URL into components for pg_dump/pg_restore.
|
|
|
|
Handles formats:
|
|
postgresql://user:pass@host:port/dbname
|
|
postgresql+asyncpg://user:pass@host:port/dbname
|
|
"""
|
|
# Remove async driver prefix if present
|
|
if "+asyncpg" in url:
|
|
url = url.replace("+asyncpg", "")
|
|
if "+psycopg2" in url:
|
|
url = url.replace("+psycopg2", "")
|
|
|
|
# Parse the URL manually to avoid dependency on urllib parsing quirks
|
|
# Format: postgresql://user:pass@host:port/dbname
|
|
rest = url.split("://", 1)[1] if "://" in url else url
|
|
user_info, rest = rest.split("@", 1) if "@" in rest else ("", rest)
|
|
user = ""
|
|
password = ""
|
|
if ":" in user_info:
|
|
user, password = user_info.split(":", 1)
|
|
else:
|
|
user = user_info
|
|
|
|
host_port, dbname = rest.split("/", 1) if "/" in rest else (rest, "")
|
|
host = host_port
|
|
port = "5432"
|
|
if ":" in host_port:
|
|
host, port = host_port.split(":", 1)
|
|
|
|
return {
|
|
"host": host,
|
|
"port": port,
|
|
"user": user,
|
|
"password": password,
|
|
"dbname": dbname,
|
|
}
|
|
|
|
|
|
async def list_backups(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
) -> list[Backup]:
|
|
"""List all backups for a tenant, ordered by creation date descending."""
|
|
stmt = (
|
|
select(Backup)
|
|
.where(Backup.tenant_id == tenant_id)
|
|
.order_by(Backup.created_at.desc())
|
|
)
|
|
result = await db.execute(stmt)
|
|
return list(result.scalars().all())
|
|
|
|
|
|
async def create_backup(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
user_id: uuid.UUID | None = None,
|
|
) -> Backup:
|
|
"""Create a database backup using pg_dump.
|
|
|
|
Creates a backup record, runs pg_dump to a file, then updates the record
|
|
with the file size and status.
|
|
"""
|
|
_ensure_backup_dir()
|
|
|
|
timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
|
|
filename = f"leocrm_backup_{tenant_id}_{timestamp}.dump"
|
|
filepath = BACKUP_DIR / filename
|
|
|
|
# Create initial pending record
|
|
backup = Backup(
|
|
tenant_id=tenant_id,
|
|
filename=filename,
|
|
status="pending",
|
|
created_by=user_id,
|
|
)
|
|
db.add(backup)
|
|
await db.flush()
|
|
await db.refresh(backup)
|
|
|
|
try:
|
|
database_url = _get_database_url()
|
|
pg = _parse_pg_url(database_url)
|
|
|
|
# Build pg_dump command
|
|
env = os.environ.copy()
|
|
if pg["password"]:
|
|
env["PGPASSWORD"] = pg["password"]
|
|
|
|
cmd = [
|
|
"pg_dump",
|
|
"--host", pg["host"],
|
|
"--port", pg["port"],
|
|
"--username", pg["user"],
|
|
"--format", "custom",
|
|
"--file", str(filepath),
|
|
pg["dbname"],
|
|
]
|
|
|
|
logger.info("Running pg_dump: %s to %s", cmd, filepath)
|
|
|
|
process = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
env=env,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
stdout, stderr = await process.communicate()
|
|
|
|
if process.returncode != 0:
|
|
error_msg = stderr.decode() if stderr else "pg_dump failed with unknown error"
|
|
logger.error("pg_dump failed: %s", error_msg)
|
|
backup.status = "failed"
|
|
backup.error_message = error_msg
|
|
await db.flush()
|
|
await db.refresh(backup)
|
|
return backup
|
|
|
|
# Get file size
|
|
size_bytes = filepath.stat().st_size if filepath.exists() else 0
|
|
|
|
# Update backup record
|
|
backup.status = "completed"
|
|
backup.size_bytes = size_bytes
|
|
backup.completed_at = datetime.now(UTC)
|
|
await db.flush()
|
|
await db.refresh(backup)
|
|
|
|
logger.info("Backup completed: %s (%d bytes)", filename, size_bytes)
|
|
return backup
|
|
|
|
except Exception as exc:
|
|
logger.exception("Backup creation failed")
|
|
backup.status = "failed"
|
|
backup.error_message = str(exc)
|
|
await db.flush()
|
|
await db.refresh(backup)
|
|
return backup
|
|
|
|
|
|
async def restore_backup(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
backup_id: uuid.UUID,
|
|
) -> Backup:
|
|
"""Restore a database backup using pg_restore.
|
|
|
|
WARNING: This is a destructive operation. It drops and recreates the database.
|
|
"""
|
|
stmt = select(Backup).where(
|
|
Backup.id == backup_id,
|
|
Backup.tenant_id == tenant_id,
|
|
)
|
|
result = await db.execute(stmt)
|
|
backup = result.scalar_one_or_none()
|
|
|
|
if backup is None:
|
|
raise ValueError("Backup not found")
|
|
|
|
if backup.status != "completed":
|
|
raise ValueError(f"Backup status is '{backup.status}', cannot restore")
|
|
|
|
filepath = BACKUP_DIR / backup.filename
|
|
if not filepath.exists():
|
|
raise FileNotFoundError(f"Backup file not found: {filepath}")
|
|
|
|
try:
|
|
database_url = _get_database_url()
|
|
pg = _parse_pg_url(database_url)
|
|
|
|
env = os.environ.copy()
|
|
if pg["password"]:
|
|
env["PGPASSWORD"] = pg["password"]
|
|
|
|
cmd = [
|
|
"pg_restore",
|
|
"--host", pg["host"],
|
|
"--port", pg["port"],
|
|
"--username", pg["user"],
|
|
"--dbname", pg["dbname"],
|
|
"--clean",
|
|
"--if-exists",
|
|
"--no-owner",
|
|
"--no-acl",
|
|
str(filepath),
|
|
]
|
|
|
|
logger.info("Running pg_restore: %s", cmd)
|
|
|
|
# Run pg_restore in a subprocess — atomic at the DB level via pg_restore --clean
|
|
import subprocess
|
|
result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=300) # noqa: ASYNC221
|
|
if result.returncode != 0:
|
|
logger.error("pg_restore failed: %s", result.stderr)
|
|
backup.status = "failed"
|
|
backup.error_message = result.stderr[:500]
|
|
await db.commit()
|
|
raise RuntimeError(f"pg_restore failed: {result.stderr[:200]}")
|
|
|
|
backup.status = "restored"
|
|
backup.restored_at = datetime.now(UTC)
|
|
await db.commit()
|
|
logger.info("Backup %s restored successfully", backup_id)
|
|
return backup
|
|
|
|
except Exception as e:
|
|
logger.error("Restore failed: %s", e)
|
|
backup.status = "failed"
|
|
backup.error_message = str(e)[:500]
|
|
await db.commit()
|
|
raise
|
|
|
|
|
|
async def delete_backup(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
backup_id: uuid.UUID,
|
|
) -> bool:
|
|
"""Delete a backup record and its file."""
|
|
stmt = select(Backup).where(
|
|
Backup.id == backup_id,
|
|
Backup.tenant_id == tenant_id,
|
|
)
|
|
result = await db.execute(stmt)
|
|
backup = result.scalar_one_or_none()
|
|
|
|
if backup is None:
|
|
return False
|
|
|
|
# Delete the file if it exists
|
|
filepath = BACKUP_DIR / backup.filename
|
|
if filepath.exists():
|
|
filepath.unlink()
|
|
|
|
await db.delete(backup)
|
|
await db.flush()
|
|
return True
|