413 lines
16 KiB
Python
413 lines
16 KiB
Python
"""Consistent PostgreSQL/S3 backup and empty-instance restore operations."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from datetime import datetime, timezone
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
import tempfile
|
|
from typing import Any
|
|
|
|
from sqlalchemy import text
|
|
|
|
from .database import SCHEMA
|
|
|
|
|
|
MANIFEST_SCHEMA = 1
|
|
VAULT_ID = re.compile(r"[0-9a-f]{32}")
|
|
CONTENT_HASH = re.compile(r"[0-9a-f]{64}")
|
|
TABLES: dict[str, tuple[str, ...]] = {
|
|
"schema_version": ("version",),
|
|
"users": ("id", "username", "password"),
|
|
"devices": ("id", "user_id", "name", "revoked"),
|
|
"sessions": ("token", "refresh", "device_id", "expires", "refresh_expires"),
|
|
"vaults": ("id", "user_id", "name", "sequence", "quota", "used"),
|
|
"uploads": ("id", "vault_id", "device_id", "hash", "size", "offset_bytes", "expires"),
|
|
"objects": ("vault_id", "hash", "size", "created"),
|
|
"revisions": (
|
|
"vault_id",
|
|
"sequence",
|
|
"file_id",
|
|
"base_revision",
|
|
"path",
|
|
"path_key",
|
|
"operation",
|
|
"hash",
|
|
"size",
|
|
"device_id",
|
|
"operation_id",
|
|
"fingerprint",
|
|
),
|
|
"files": ("vault_id", "file_id", "sequence", "path_key", "deleted"),
|
|
"login_limits": ("key", "started", "attempts"),
|
|
"upload_receipts": ("id", "vault_id", "device_id", "hash", "completed"),
|
|
}
|
|
|
|
|
|
class OperationsError(RuntimeError):
|
|
"""Stable operator-facing failure without credentials or response bodies."""
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _object_path(root: Path, vault_id: str, content_hash: str) -> Path:
|
|
if not VAULT_ID.fullmatch(vault_id) or not CONTENT_HASH.fullmatch(content_hash):
|
|
raise OperationsError("OBJECT_ID_INVALID")
|
|
return root / "objects" / vault_id / content_hash
|
|
|
|
|
|
def _manifest_objects(conn) -> list[dict[str, Any]]:
|
|
pending = conn.execute(text("SELECT COUNT(*) FROM uploads")).scalar_one()
|
|
if pending:
|
|
raise OperationsError("BACKUP_PENDING_UPLOADS")
|
|
missing = conn.execute(
|
|
text(
|
|
"SELECT COUNT(*) FROM revisions r LEFT JOIN objects o "
|
|
"ON o.vault_id=r.vault_id AND o.hash=r.hash "
|
|
"WHERE r.hash IS NOT NULL AND o.hash IS NULL"
|
|
)
|
|
).scalar_one()
|
|
if missing:
|
|
raise OperationsError("HISTORICAL_OBJECT_MISSING")
|
|
return [
|
|
{"vault_id": row.vault_id, "hash": row.hash, "size": int(row.size)}
|
|
for row in conn.execute(
|
|
text("SELECT vault_id,hash,size FROM objects ORDER BY vault_id,hash")
|
|
)
|
|
]
|
|
|
|
|
|
def _write_database_snapshot(conn, destination: Path) -> dict[str, int]:
|
|
counts = {}
|
|
with destination.open("x", encoding="utf-8", newline="\n") as output:
|
|
for table, columns in TABLES.items():
|
|
projection = ",".join(f'"{column}"' for column in columns)
|
|
ordering = ",".join(f'"{column}"' for column in columns)
|
|
rows = conn.execute(
|
|
text(f'SELECT {projection} FROM "{table}" ORDER BY {ordering}')
|
|
)
|
|
count = 0
|
|
for row in rows:
|
|
output.write(
|
|
json.dumps(
|
|
{"table": table, "values": list(row)},
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
)
|
|
+ "\n"
|
|
)
|
|
count += 1
|
|
counts[table] = count
|
|
output.flush()
|
|
os.fsync(output.fileno())
|
|
try:
|
|
destination.chmod(0o600)
|
|
except OSError:
|
|
pass
|
|
return counts
|
|
|
|
|
|
def _download(objects, root: Path, item: dict[str, Any]) -> None:
|
|
target = _object_path(root, item["vault_id"], item["hash"])
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
response = objects.client.get_object(
|
|
Bucket=objects.bucket,
|
|
Key=f"{item['vault_id']}/{item['hash']}",
|
|
)
|
|
digest = hashlib.sha256()
|
|
size = 0
|
|
try:
|
|
with response["Body"] as body, target.open("xb") as destination:
|
|
for chunk in iter(lambda: body.read(1024 * 1024), b""):
|
|
destination.write(chunk)
|
|
digest.update(chunk)
|
|
size += len(chunk)
|
|
destination.flush()
|
|
os.fsync(destination.fileno())
|
|
except BaseException:
|
|
target.unlink(missing_ok=True)
|
|
raise
|
|
if size != item["size"] or digest.hexdigest() != item["hash"]:
|
|
target.unlink(missing_ok=True)
|
|
raise OperationsError("OBJECT_INTEGRITY_FAILED")
|
|
|
|
|
|
def create_backup(db, objects, destination: Path, *, workers: int = 8) -> dict[str, Any]:
|
|
destination = destination.resolve(strict=False)
|
|
if destination.exists():
|
|
raise OperationsError("BACKUP_DESTINATION_EXISTS")
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = Path(
|
|
tempfile.mkdtemp(prefix=destination.name + ".incomplete-", dir=destination.parent)
|
|
)
|
|
database = temporary / "database.jsonl"
|
|
try:
|
|
with db.engine.connect() as conn:
|
|
transaction = conn.begin()
|
|
try:
|
|
conn.exec_driver_sql(
|
|
"SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY"
|
|
)
|
|
version = conn.execute(text("SELECT version FROM schema_version")).scalar_one()
|
|
if version != 1:
|
|
raise OperationsError("SCHEMA_INCOMPATIBLE")
|
|
catalog = _manifest_objects(conn)
|
|
table_rows = _write_database_snapshot(conn, database)
|
|
transaction.commit()
|
|
except BaseException:
|
|
if transaction.is_active:
|
|
transaction.rollback()
|
|
raise
|
|
with ThreadPoolExecutor(max_workers=workers) as pool:
|
|
list(pool.map(lambda item: _download(objects, temporary, item), catalog))
|
|
manifest = {
|
|
"schema": MANIFEST_SCHEMA,
|
|
"created_utc": datetime.now(timezone.utc).isoformat(),
|
|
"database_sha256": sha256_file(database),
|
|
"database_rows": table_rows,
|
|
"object_count": len(catalog),
|
|
"object_bytes": sum(item["size"] for item in catalog),
|
|
"objects": catalog,
|
|
}
|
|
manifest_path = temporary / "manifest.json"
|
|
manifest_path.write_text(
|
|
json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
os.replace(temporary, destination)
|
|
return {
|
|
"status": "BACKUP_COMPLETE",
|
|
"created_utc": manifest["created_utc"],
|
|
"object_count": manifest["object_count"],
|
|
"object_bytes": manifest["object_bytes"],
|
|
}
|
|
except BaseException:
|
|
shutil.rmtree(temporary, ignore_errors=True)
|
|
raise
|
|
|
|
|
|
def _load_manifest(source: Path, max_age_hours: float) -> dict[str, Any]:
|
|
try:
|
|
manifest = json.loads((source / "manifest.json").read_text(encoding="utf-8"))
|
|
created = datetime.fromisoformat(manifest["created_utc"])
|
|
except (OSError, ValueError, KeyError, json.JSONDecodeError) as error:
|
|
raise OperationsError("BACKUP_MANIFEST_INVALID") from error
|
|
if manifest.get("schema") != MANIFEST_SCHEMA or created.tzinfo is None:
|
|
raise OperationsError("BACKUP_MANIFEST_INVALID")
|
|
age = (datetime.now(timezone.utc) - created.astimezone(timezone.utc)).total_seconds()
|
|
if age < -300 or age > max_age_hours * 3600:
|
|
raise OperationsError("BACKUP_AGE_INVALID")
|
|
catalog = manifest.get("objects")
|
|
if not isinstance(catalog, list) or manifest.get("database_rows") is None:
|
|
raise OperationsError("BACKUP_MANIFEST_INVALID")
|
|
normalized = []
|
|
for item in catalog:
|
|
if (
|
|
not isinstance(item, dict)
|
|
or not isinstance(item.get("vault_id"), str)
|
|
or not isinstance(item.get("hash"), str)
|
|
or not isinstance(item.get("size"), int)
|
|
or isinstance(item.get("size"), bool)
|
|
or item["size"] < 0
|
|
):
|
|
raise OperationsError("BACKUP_MANIFEST_INVALID")
|
|
_object_path(source, item["vault_id"], item["hash"])
|
|
normalized.append(
|
|
{"vault_id": item["vault_id"], "hash": item["hash"], "size": item["size"]}
|
|
)
|
|
rows = manifest["database_rows"]
|
|
if (
|
|
normalized != sorted(normalized, key=lambda item: (item["vault_id"], item["hash"]))
|
|
or len({(item["vault_id"], item["hash"]) for item in normalized}) != len(normalized)
|
|
or manifest.get("object_count") != len(normalized)
|
|
or manifest.get("object_bytes") != sum(item["size"] for item in normalized)
|
|
or not CONTENT_HASH.fullmatch(str(manifest.get("database_sha256", "")))
|
|
or not isinstance(rows, dict)
|
|
or set(rows) != set(TABLES)
|
|
or any(not isinstance(rows[name], int) or rows[name] < 0 for name in TABLES)
|
|
):
|
|
raise OperationsError("BACKUP_MANIFEST_INVALID")
|
|
manifest["objects"] = normalized
|
|
manifest["age_seconds"] = max(0, int(age))
|
|
return manifest
|
|
|
|
|
|
def _load_database_snapshot(source: Path, manifest: dict[str, Any]) -> dict[str, list[list]]:
|
|
database = source / "database.jsonl"
|
|
if not database.is_file() or sha256_file(database) != manifest["database_sha256"]:
|
|
raise OperationsError("BACKUP_DATABASE_INTEGRITY_FAILED")
|
|
restored = {table: [] for table in TABLES}
|
|
try:
|
|
with database.open("r", encoding="utf-8") as stream:
|
|
for line in stream:
|
|
item = json.loads(line)
|
|
table = item.get("table")
|
|
values = item.get("values")
|
|
if table not in TABLES or not isinstance(values, list):
|
|
raise OperationsError("BACKUP_DATABASE_INVALID")
|
|
if len(values) != len(TABLES[table]):
|
|
raise OperationsError("BACKUP_DATABASE_INVALID")
|
|
restored[table].append(values)
|
|
except (OSError, json.JSONDecodeError) as error:
|
|
raise OperationsError("BACKUP_DATABASE_INVALID") from error
|
|
if any(len(restored[table]) != manifest["database_rows"][table] for table in TABLES):
|
|
raise OperationsError("BACKUP_DATABASE_INVALID")
|
|
if restored["schema_version"] != [[1]]:
|
|
raise OperationsError("BACKUP_DATABASE_INVALID")
|
|
object_columns = TABLES["objects"]
|
|
indexes = {name: object_columns.index(name) for name in ("vault_id", "hash", "size")}
|
|
catalog = sorted(
|
|
(
|
|
{
|
|
"vault_id": row[indexes["vault_id"]],
|
|
"hash": row[indexes["hash"]],
|
|
"size": row[indexes["size"]],
|
|
}
|
|
for row in restored["objects"]
|
|
),
|
|
key=lambda item: (item["vault_id"], item["hash"]),
|
|
)
|
|
if catalog != manifest["objects"]:
|
|
raise OperationsError("BACKUP_DATABASE_INVALID")
|
|
return restored
|
|
|
|
|
|
def _verify_backup_file(source: Path, item: dict[str, Any]) -> None:
|
|
path = _object_path(source, item["vault_id"], item["hash"])
|
|
try:
|
|
size = path.stat().st_size
|
|
except OSError as error:
|
|
raise OperationsError("BACKUP_OBJECT_MISSING") from error
|
|
if size != item["size"] or sha256_file(path) != item["hash"]:
|
|
raise OperationsError("BACKUP_OBJECT_INTEGRITY_FAILED")
|
|
|
|
|
|
def _upload_and_verify(objects, source: Path, item: dict[str, Any]) -> str:
|
|
key = f"{item['vault_id']}/{item['hash']}"
|
|
path = _object_path(source, item["vault_id"], item["hash"])
|
|
objects.put_file(key, path, item["hash"])
|
|
try:
|
|
response = objects.client.get_object(Bucket=objects.bucket, Key=key)
|
|
digest = hashlib.sha256()
|
|
size = 0
|
|
with response["Body"] as body:
|
|
for chunk in iter(lambda: body.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
size += len(chunk)
|
|
if size != item["size"] or digest.hexdigest() != item["hash"]:
|
|
raise OperationsError("RESTORED_OBJECT_INTEGRITY_FAILED")
|
|
except BaseException:
|
|
try:
|
|
objects.delete(key)
|
|
except BaseException:
|
|
pass
|
|
raise
|
|
return key
|
|
|
|
|
|
def _database_is_empty(conn) -> bool:
|
|
return (
|
|
conn.execute(
|
|
text(
|
|
"SELECT COUNT(*) FROM information_schema.tables "
|
|
"WHERE table_schema='public'"
|
|
)
|
|
).scalar_one()
|
|
== 0
|
|
)
|
|
|
|
|
|
def _restore_database(
|
|
db, snapshot: dict[str, list[list]], expected_objects: list[dict[str, Any]]
|
|
) -> None:
|
|
with db.transaction() as conn:
|
|
conn.execute(text("SELECT pg_advisory_xact_lock(1330534488)"))
|
|
if not _database_is_empty(conn):
|
|
raise OperationsError("RESTORE_DATABASE_NOT_EMPTY")
|
|
for statement in SCHEMA:
|
|
conn.execute(text(statement))
|
|
for table, columns in TABLES.items():
|
|
rows = snapshot[table]
|
|
if not rows:
|
|
continue
|
|
names = ",".join(f'"{column}"' for column in columns)
|
|
values = ",".join(f":v{index}" for index in range(len(columns)))
|
|
parameters = [
|
|
{f"v{index}": value for index, value in enumerate(row)} for row in rows
|
|
]
|
|
conn.execute(text(f'INSERT INTO "{table}" ({names}) VALUES ({values})'), parameters)
|
|
version = conn.execute(text("SELECT version FROM schema_version")).scalar_one()
|
|
if version != 1 or _manifest_objects(conn) != expected_objects:
|
|
raise OperationsError("RESTORED_DATABASE_INTEGRITY_FAILED")
|
|
|
|
|
|
def restore_backup(
|
|
db,
|
|
objects,
|
|
source: Path,
|
|
*,
|
|
workers: int = 8,
|
|
max_age_hours: float = 24,
|
|
) -> dict[str, Any]:
|
|
if not math.isfinite(max_age_hours) or max_age_hours <= 0:
|
|
raise OperationsError("BACKUP_AGE_LIMIT_INVALID")
|
|
source = source.resolve(strict=True)
|
|
manifest = _load_manifest(source, max_age_hours)
|
|
snapshot = _load_database_snapshot(source, manifest)
|
|
with ThreadPoolExecutor(max_workers=workers) as pool:
|
|
list(pool.map(lambda item: _verify_backup_file(source, item), manifest["objects"]))
|
|
with db.engine.connect() as conn:
|
|
if not _database_is_empty(conn):
|
|
raise OperationsError("RESTORE_DATABASE_NOT_EMPTY")
|
|
created_bucket = objects.ensure_bucket()
|
|
if not objects.is_empty():
|
|
raise OperationsError("RESTORE_BUCKET_NOT_EMPTY")
|
|
uploaded: list[str] = []
|
|
database_restored = False
|
|
try:
|
|
failure = None
|
|
with ThreadPoolExecutor(max_workers=workers) as pool:
|
|
futures = [
|
|
pool.submit(_upload_and_verify, objects, source, item)
|
|
for item in manifest["objects"]
|
|
]
|
|
for future in as_completed(futures):
|
|
try:
|
|
uploaded.append(future.result())
|
|
except BaseException as error:
|
|
failure = failure or error
|
|
if failure is not None:
|
|
raise failure
|
|
_restore_database(db, snapshot, manifest["objects"])
|
|
database_restored = True
|
|
except BaseException:
|
|
if not database_restored:
|
|
try:
|
|
objects.delete_many(uploaded)
|
|
except BaseException:
|
|
pass
|
|
if created_bucket:
|
|
try:
|
|
objects.delete_bucket()
|
|
except BaseException:
|
|
pass
|
|
raise
|
|
return {
|
|
"status": "RESTORE_COMPLETE",
|
|
"backup_age_seconds": manifest["age_seconds"],
|
|
"object_count": manifest["object_count"],
|
|
"object_bytes": manifest["object_bytes"],
|
|
"verified_objects": len(manifest["objects"]),
|
|
}
|