Files
NotesAgentic/server sync/sync_server/operations.py
T
admin d703ab64e3
CI / docs-check (push) Canceled after 0s
CI / backend-test (push) Canceled after 0s
CI / service-test (push) Canceled after 0s
CI / frontend-test (push) Canceled after 0s
CI / rust-core (push) Canceled after 0s
CI / docs-check (pull_request) Canceled after 0s
CI / backend-test (pull_request) Canceled after 0s
CI / service-test (pull_request) Canceled after 0s
CI / frontend-test (pull_request) Canceled after 0s
CI / rust-core (pull_request) Canceled after 0s
docs: 将仓库代码注释统一为中文
2026-09-10 00:40:56 +08:00

413 lines
16 KiB
Python

"""一致的 PostgreSQL/S3 备份和空实例恢复操作。"""
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):
"""稳定的面向操作员的故障,无需凭证或响应主体。"""
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"]),
}