Files
sales-trainer/backend/scripts/migrate_json_to_postgres.py
Macky 2c40fc7502 feat(import): commit S4.4 JSON→PostgreSQL importer + API error-handler hardening
Re-verified staged increment from a clean requirements.lock.txt venv:
- 330 backend tests pass (17/17 in new error_handlers + json_import tests)
- compileall + frontend npm build clean
- git diff --check clean; no secrets in diff
- importer CLI dry-run bootstrap works

Includes JSON HTTPException handler under /api/* and parse-safe static 404
via abort. JSON stores remain runtime-authoritative; production operation
still gated behind operator approval.
2026-08-16 07:42:32 +07:00

696 lines
28 KiB
Python

"""Fail-closed, idempotent JSON-store to relational importer.
Dry-run is the default. Applying data requires an explicit target URL and a
fresh backup directory. The importer validates the complete reference graph
before opening a target transaction and never overwrites conflicting rows.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import re
import shutil
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
# Allow the documented ``python scripts/migrate_json_to_postgres.py`` form to
# resolve the backend's ``app`` package when this file is executed directly.
if __package__ in (None, ""):
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sqlalchemy.exc import SQLAlchemyError
from app.db import create_db_engine, create_session_factory
from app.models import Group, Message, Organization, Persona, TrainingSession, User
from app.services.store import revealable_view
MAX_SOURCE_FILE_BYTES = 10 * 1024 * 1024
MAX_SOURCE_TOTAL_BYTES = 64 * 1024 * 1024
MAX_SOURCE_FILES = 5_000
MAX_SOURCE_RECORDS = 50_000
MAX_SOURCE_ROWS = 100_000
MAX_SOURCE_MESSAGES = 100_000
MAX_SOURCE_JSON_DEPTH = 32
_ID_RE = re.compile(r"^[A-Za-z0-9_.:-]{1,64}$")
_COLLECTIONS = ("orgs", "users", "groups", "sessions", "my_personas")
class ImportValidationError(ValueError):
"""Raised when source data or a target conflict fails closed."""
@dataclass(frozen=True)
class SourceSnapshot:
orgs: list[dict[str, Any]]
users: list[dict[str, Any]]
groups: list[dict[str, Any]]
sessions: list[dict[str, Any]]
my_personas: list[dict[str, Any]]
@dataclass
class _SourceBudget:
files: int = 0
bytes: int = 0
records: int = 0
def account_file(self, size: int) -> None:
self.files += 1
self.bytes += size
if self.files > MAX_SOURCE_FILES:
raise ImportValidationError("source file limit exceeded")
if self.bytes > MAX_SOURCE_TOTAL_BYTES:
raise ImportValidationError("total source size exceeds bounded limit")
def account_record(self) -> None:
self.records += 1
if self.records > MAX_SOURCE_RECORDS:
raise ImportValidationError("source record limit exceeded")
@dataclass(frozen=True)
class ImportRow:
table: str
key: str
model: type[Any]
values: dict[str, Any]
def _validate_json_depth(value: object, collection: str) -> None:
stack: list[tuple[object, int]] = [(value, 1)]
while stack:
current, depth = stack.pop()
if depth > MAX_SOURCE_JSON_DEPTH:
raise ImportValidationError(f"source JSON nesting limit exceeded: {collection}")
if isinstance(current, dict):
stack.extend((child, depth + 1) for child in current.values())
elif isinstance(current, list):
stack.extend((child, depth + 1) for child in current)
def _read_bounded_bytes(path: Path, *, max_bytes: int, too_large_message: str) -> bytes:
with path.open("rb") as handle:
content = handle.read(max_bytes + 1)
if len(content) > max_bytes:
raise ImportValidationError(too_large_message)
return content
def _read_collection(
source_root: Path,
collection: str,
budget: _SourceBudget,
) -> list[dict[str, Any]]:
directory = source_root / collection
if not directory.exists():
return []
if directory.is_symlink() or not directory.is_dir():
raise ImportValidationError(f"source collection is not a directory: {collection}")
paths: list[Path] = []
try:
for path in directory.iterdir():
if path.suffix != ".json":
continue
if path.is_symlink():
raise ImportValidationError(f"symlink source file rejected: {collection}")
size = path.stat().st_size
if size > MAX_SOURCE_FILE_BYTES:
raise ImportValidationError(f"source file exceeds bounded size: {collection}")
budget.account_file(size)
paths.append(path)
except ImportValidationError:
raise
except OSError as exc:
raise ImportValidationError(f"source collection cannot be read: {collection}") from exc
rows: list[dict[str, Any]] = []
for path in sorted(paths, key=lambda item: item.name):
try:
raw = _read_bounded_bytes(
path,
max_bytes=MAX_SOURCE_FILE_BYTES,
too_large_message=f"source file exceeds bounded size: {collection}",
)
value = json.loads(raw.decode("utf-8"))
except ImportValidationError:
raise
except RecursionError as exc:
raise ImportValidationError(f"source JSON nesting limit exceeded: {collection}") from exc
except (OSError, UnicodeError, ValueError) as exc:
raise ImportValidationError(f"invalid source JSON in {collection}") from exc
if not isinstance(value, dict):
raise ImportValidationError(f"source record must be an object: {collection}")
_validate_json_depth(value, collection)
budget.account_record()
rows.append(value)
return rows
def _reject_nonempty_audit(source_root: Path, budget: _SourceBudget) -> None:
path = source_root / "audit" / "audit.jsonl"
if not path.exists():
return
if path.is_symlink():
raise ImportValidationError("symlink audit source rejected")
try:
size = path.stat().st_size
if size > MAX_SOURCE_FILE_BYTES:
raise ImportValidationError("audit source exceeds bounded size")
budget.account_file(size)
raw = _read_bounded_bytes(
path,
max_bytes=MAX_SOURCE_FILE_BYTES,
too_large_message="audit source exceeds bounded size",
)
has_entries = any(line.strip() for line in raw.decode("utf-8").splitlines())
except ImportValidationError:
raise
except (OSError, UnicodeError) as exc:
raise ImportValidationError("audit source cannot be read") from exc
if has_entries:
raise ImportValidationError(
"audit JSONL contains entries; import audit events through the S4.6 PostgreSQL audit gate"
)
def load_source(source_root: Path) -> SourceSnapshot:
"""Read only supported JSON collections; credentials/config files are ignored."""
root = source_root.expanduser().resolve()
if not root.exists() or not root.is_dir():
raise ImportValidationError("source data directory is missing")
budget = _SourceBudget()
_reject_nonempty_audit(root, budget)
values = {
collection: _read_collection(root, collection, budget)
for collection in _COLLECTIONS
}
return SourceSnapshot(**values)
def _canonical_checksum(snapshot: SourceSnapshot) -> str:
payload = {collection: getattr(snapshot, collection) for collection in _COLLECTIONS}
encoded = json.dumps(
payload,
ensure_ascii=False,
sort_keys=True,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(encoded).hexdigest()
def _record_id(record: dict[str, Any], collection: str) -> str:
value = record.get("id")
if not isinstance(value, str) or not _ID_RE.fullmatch(value):
raise ImportValidationError(f"invalid id in {collection} collection")
return value
def _required_string(value: object, field: str, collection: str) -> str:
if not isinstance(value, str) or not value.strip():
raise ImportValidationError(f"invalid {field} in {collection} collection")
return value.strip()
def _optional_string(value: object, field: str, collection: str) -> str | None:
if value is None or value == "":
return None
return _required_string(value, field, collection)
def _timestamp(value: object, field: str, collection: str) -> datetime | None:
if value is None or value == "":
return None
if not isinstance(value, str):
raise ImportValidationError(f"invalid {field} timestamp in {collection} collection")
try:
parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
except ValueError as exc:
raise ImportValidationError(f"invalid {field} timestamp in {collection} collection") from exc
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=timezone.utc)
return parsed.astimezone(timezone.utc)
def _json_object(value: object, field: str, collection: str) -> dict[str, Any] | None:
if value is None:
return None
if not isinstance(value, dict):
raise ImportValidationError(f"invalid {field} object in {collection} collection")
return value
def _bool_value(value: object, field: str, collection: str, *, default: bool) -> bool:
if value is None:
return default
if not isinstance(value, bool):
raise ImportValidationError(f"invalid {field} in {collection} collection")
return value
def _int_value(value: object, field: str, collection: str, *, default: int, minimum: int = 0) -> int:
if value is None:
return default
if isinstance(value, bool) or not isinstance(value, int) or value < minimum:
raise ImportValidationError(f"invalid {field} in {collection} collection")
return value
def _validate_unique(values: list[str], label: str) -> None:
if len(values) != len(set(values)):
raise ImportValidationError(f"duplicate {label} in source")
def _scenario_payload(row: dict[str, Any]) -> dict[str, Any] | None:
raw = row.get("scenario")
if isinstance(raw, dict):
payload = dict(raw)
if row.get("locale") is not None:
payload.setdefault("locale", row.get("locale"))
if row.get("persona_name") is not None:
payload.setdefault("persona_name", row.get("persona_name"))
if row.get("persona_meta") is not None:
payload.setdefault("persona_meta", row.get("persona_meta"))
return payload
if raw is None and not any(row.get(key) is not None for key in ("locale", "persona_name", "persona_meta")):
return None
return {
"scenario": raw,
"locale": row.get("locale"),
"persona_name": row.get("persona_name"),
"persona_meta": row.get("persona_meta") or {},
}
def _report_text(value: object) -> str | None:
if value is None:
return None
if isinstance(value, str):
return value
return json.dumps(value, ensure_ascii=False, sort_keys=True)
def _derived_message_id(session_id: str, sequence: int) -> str:
digest = hashlib.sha256(f"{session_id}\0{sequence}".encode("utf-8")).hexdigest()[:32]
return f"msg-{digest}"
def _append_row(rows: list[ImportRow], row: ImportRow) -> None:
if len(rows) >= MAX_SOURCE_ROWS:
raise ImportValidationError("source row limit exceeded")
rows.append(row)
def _queue_persona_record(
rows: list[ImportRow],
persona_records: list[tuple[str, str, dict[str, Any], str]],
record: tuple[str, str, dict[str, Any], str],
) -> None:
if len(rows) + len(persona_records) >= MAX_SOURCE_ROWS:
raise ImportValidationError("source row limit exceeded")
persona_records.append(record)
def _build_rows(snapshot: SourceSnapshot) -> list[ImportRow]:
orgs_by_id: dict[str, dict[str, Any]] = {}
users_by_id: dict[str, dict[str, Any]] = {}
groups_by_id: dict[str, dict[str, Any]] = {}
personas_by_id: dict[str, tuple[str, dict[str, Any]]] = {}
rows: list[ImportRow] = []
for record in snapshot.orgs:
org_id = _record_id(record, "orgs")
if org_id in orgs_by_id:
raise ImportValidationError("duplicate organization id in source")
orgs_by_id[org_id] = record
values = {
"id": org_id,
"name": _required_string(record.get("name"), "name", "orgs"),
"plan": _required_string(record.get("plan", "trial"), "plan", "orgs"),
"seats": _int_value(record.get("seats"), "seats", "orgs", default=5, minimum=1),
"active": _bool_value(record.get("active"), "active", "orgs", default=True),
}
created_at = _timestamp(record.get("created_at"), "created_at", "orgs")
if created_at is not None:
values["created_at"] = created_at
_append_row(rows, ImportRow("organizations", org_id, Organization, values))
usernames: list[str] = []
emails: list[str] = []
for record in snapshot.users:
user_id = _record_id(record, "users")
if user_id in users_by_id:
raise ImportValidationError("duplicate user id in source")
org_id = _required_string(record.get("org_id"), "org_id", "users")
if org_id not in orgs_by_id:
raise ImportValidationError("user references missing organization")
username = _required_string(record.get("username"), "username", "users").lower()
password_hash = _required_string(record.get("password_hash"), "password_hash", "users")
email = _optional_string(record.get("email"), "email", "users")
usernames.append(username)
if email:
emails.append(email.lower())
users_by_id[user_id] = record
values = {
"id": user_id,
"org_id": org_id,
"username": username,
"password_hash": password_hash,
"email": email.lower() if email else None,
"name": _required_string(record.get("name") or username, "name", "users"),
"role": _required_string(record.get("role", "user"), "role", "users"),
"active": _bool_value(record.get("active"), "active", "users", default=True),
"must_setup": _bool_value(record.get("must_setup"), "must_setup", "users", default=False),
"accepted_terms": _bool_value(record.get("accepted_terms"), "accepted_terms", "users", default=False),
"auth_version": _int_value(record.get("auth_version"), "auth_version", "users", default=0),
}
accepted_at = _timestamp(record.get("accepted_terms_at"), "accepted_terms_at", "users")
created_at = _timestamp(record.get("created_at"), "created_at", "users")
if accepted_at is not None:
values["accepted_terms_at"] = accepted_at
if created_at is not None:
values["created_at"] = created_at
_append_row(rows, ImportRow("users", user_id, User, values))
_validate_unique(usernames, "username")
_validate_unique(emails, "email")
persona_records: list[tuple[str, str, dict[str, Any], str]] = []
for record in snapshot.groups:
group_id = _record_id(record, "groups")
if group_id in groups_by_id:
raise ImportValidationError("duplicate group id in source")
org_id = _required_string(record.get("org_id"), "org_id", "groups")
if org_id not in orgs_by_id:
raise ImportValidationError("cross-tenant group reference: missing organization")
creator_id = _optional_string(
record.get("creator_user_id") or record.get("creator_id"),
"creator_user_id",
"groups",
)
owner_id = _optional_string(record.get("owner_user_id"), "owner_user_id", "groups")
for label, user_id in (("creator", creator_id), ("owner", owner_id)):
if user_id is None:
continue
user = users_by_id.get(user_id)
if user is None or user.get("org_id") != org_id:
raise ImportValidationError(f"cross-tenant group reference: {label}")
groups_by_id[group_id] = record
values = {
"id": group_id,
"org_id": org_id,
"owner_user_id": owner_id,
"creator_user_id": creator_id,
"name": _required_string(record.get("name") or record.get("title"), "name", "groups"),
"status": _required_string(record.get("status", "draft"), "status", "groups"),
"input_json": _json_object(record.get("input"), "input", "groups"),
"sales_kit_json": _json_object(record.get("sales_kit"), "sales_kit", "groups"),
"report_text": _report_text(record.get("report")),
}
created_at = _timestamp(record.get("created_at"), "created_at", "groups")
updated_at = _timestamp(record.get("updated_at"), "updated_at", "groups")
if created_at is not None:
values["created_at"] = created_at
if updated_at is not None:
values["updated_at"] = updated_at
_append_row(rows, ImportRow("groups", group_id, Group, values))
personas = record.get("personas", [])
if not isinstance(personas, list):
raise ImportValidationError("group personas must be a list")
for persona in personas:
if not isinstance(persona, dict):
raise ImportValidationError("persona record must be an object")
_queue_persona_record(rows, persona_records, (group_id, org_id, persona, "groups"))
for record in snapshot.my_personas:
user_id = _required_string(record.get("user_id"), "user_id", "my_personas")
user = users_by_id.get(user_id)
if user is None:
raise ImportValidationError("my_personas references missing user")
persona = record.get("persona")
if not isinstance(persona, dict):
raise ImportValidationError("my_personas persona must be an object")
private_group_id = "private-" + hashlib.sha256(user_id.encode("utf-8")).hexdigest()[:24]
existing_private_group = groups_by_id.get(private_group_id)
if existing_private_group is not None:
existing_creator_id = existing_private_group.get("creator_user_id") or existing_private_group.get("creator_id")
if (
existing_private_group.get("org_id") != user["org_id"]
or existing_private_group.get("owner_user_id") != user_id
or existing_creator_id != user_id
or existing_private_group.get("status") != "ready"
):
raise ImportValidationError("private group id collision")
else:
groups_by_id[private_group_id] = {
"id": private_group_id,
"org_id": user["org_id"],
"owner_user_id": user_id,
"creator_user_id": user_id,
"status": "ready",
}
_append_row(rows,
ImportRow(
"groups",
private_group_id,
Group,
{
"id": private_group_id,
"org_id": user["org_id"],
"owner_user_id": user_id,
"creator_user_id": user_id,
"name": f"{user.get('name') or user_id}'s private personas",
"status": "ready",
"input_json": {"source": "my_personas"},
"sales_kit_json": None,
"report_text": None,
},
)
)
_queue_persona_record(rows, persona_records, (private_group_id, user["org_id"], persona, "my_personas"))
for group_id, _org_id, persona, collection in persona_records:
persona_id = _record_id(persona, collection)
if persona_id in personas_by_id:
raise ImportValidationError("duplicate persona id in source")
personas_by_id[persona_id] = (group_id, persona)
full_persona = dict(persona)
values = {
"id": persona_id,
"group_id": group_id,
"source_persona_id": _optional_string(persona.get("source_persona_id"), "source_persona_id", collection),
"tier": _required_string(persona.get("tier") or persona.get("intent_tier") or "B", "tier", collection),
"public_json": revealable_view(full_persona),
"latent_json": full_persona,
}
_append_row(rows, ImportRow("personas", persona_id, Persona, values))
session_ids: set[str] = set()
message_keys: set[tuple[str, int]] = set()
message_count = 0
for record in snapshot.sessions:
session_id = _record_id(record, "sessions")
if session_id in session_ids:
raise ImportValidationError("duplicate session id in source")
session_ids.add(session_id)
org_id = _required_string(record.get("org_id"), "org_id", "sessions")
user_id = _required_string(record.get("user_id"), "user_id", "sessions")
group_id = _required_string(record.get("group_id"), "group_id", "sessions")
persona_id = _required_string(record.get("persona_id"), "persona_id", "sessions")
user = users_by_id.get(user_id)
group = groups_by_id.get(group_id)
persona = personas_by_id.get(persona_id)
if (
user is None
or group is None
or persona is None
or user.get("org_id") != org_id
or group.get("org_id", org_id) != org_id
or persona[0] != group_id
):
raise ImportValidationError("cross-tenant session reference")
mode = _required_string(record.get("mode") or "trainee", "mode", "sessions")
if mode not in ("trainee", "preview"):
raise ImportValidationError("invalid session mode")
values = {
"id": session_id,
"org_id": org_id,
"user_id": user_id,
"group_id": group_id,
"persona_id": persona_id,
"mode": mode,
"status": _required_string(record.get("status", "active"), "status", "sessions"),
"outcome": _optional_string(record.get("outcome"), "outcome", "sessions"),
"scenario_json": _scenario_payload(record),
"internal_json": _json_object(record.get("internal"), "internal", "sessions"),
"debrief_json": _json_object(record.get("debrief"), "debrief", "sessions"),
}
created_at = _timestamp(record.get("created_at"), "created_at", "sessions")
updated_at = _timestamp(record.get("updated_at"), "updated_at", "sessions")
if created_at is not None:
values["created_at"] = created_at
if updated_at is not None:
values["updated_at"] = updated_at
_append_row(rows, ImportRow("sessions", session_id, TrainingSession, values))
messages = record.get("messages", [])
if not isinstance(messages, list):
raise ImportValidationError("session messages must be a list")
for index, message in enumerate(messages, start=1):
message_count += 1
if message_count > MAX_SOURCE_MESSAGES:
raise ImportValidationError("source message limit exceeded")
if not isinstance(message, dict):
raise ImportValidationError("message record must be an object")
sequence = message.get("sequence", index)
if isinstance(sequence, bool) or not isinstance(sequence, int) or sequence < 1:
raise ImportValidationError("message sequence must be a positive integer")
if (session_id, sequence) in message_keys:
raise ImportValidationError("duplicate message sequence in source")
message_keys.add((session_id, sequence))
role = _required_string(message.get("role"), "role", "sessions")
text = message.get("text")
if not isinstance(text, str):
raise ImportValidationError("message text must be a string")
message_values = {
"id": _derived_message_id(session_id, sequence),
"session_id": session_id,
"sequence": sequence,
"role": role,
"text": text,
}
message_time = _timestamp(message.get("created_at") or message.get("ts"), "message time", "sessions")
if message_time is not None:
message_values["created_at"] = message_time
_append_row(rows, ImportRow("messages", message_values["id"], Message, message_values))
if len(rows) > MAX_SOURCE_ROWS:
raise ImportValidationError("source row limit exceeded")
return rows
def _values_equal(actual: object, expected: object) -> bool:
if isinstance(actual, datetime) and isinstance(expected, datetime):
actual_utc = actual.replace(tzinfo=timezone.utc) if actual.tzinfo is None else actual.astimezone(timezone.utc)
expected_utc = expected.replace(tzinfo=timezone.utc) if expected.tzinfo is None else expected.astimezone(timezone.utc)
return actual_utc == expected_utc
return actual == expected
def _row_matches(existing: object, values: dict[str, Any]) -> bool:
return all(_values_equal(getattr(existing, field), value) for field, value in values.items())
def _counts(rows: list[ImportRow], *, apply: bool) -> dict[str, dict[str, int]]:
output: dict[str, dict[str, int]] = {}
for row in rows:
bucket = output.setdefault(row.table, {"created": 0, "unchanged": 0, "would_create": 0})
if apply:
bucket.pop("would_create", None)
else:
bucket["would_create"] += 1
return output
def _backup_source(source_root: Path, backup_dir: Path) -> None:
source = source_root.expanduser().resolve()
backup = backup_dir.expanduser().resolve()
try:
backup.relative_to(source)
except ValueError:
pass
else:
raise ImportValidationError("backup directory cannot be inside source data")
if backup.exists():
raise ImportValidationError("backup directory already exists")
try:
backup.mkdir(parents=True, exist_ok=False)
for collection in _COLLECTIONS:
source_collection = source / collection
if source_collection.exists():
shutil.copytree(source_collection, backup / collection, symlinks=False)
except OSError as exc:
raise ImportValidationError("source backup could not be created") from exc
def _apply_rows(rows: list[ImportRow], database_url: str) -> dict[str, dict[str, int]]:
engine = create_db_engine(database_url)
factory = create_session_factory(engine)
counts = _counts(rows, apply=True)
with factory() as session:
with session.begin():
for row in rows:
existing = session.get(row.model, row.key)
bucket = counts[row.table]
if existing is not None:
if not _row_matches(existing, row.values):
raise ImportValidationError(f"target conflict for {row.table} record")
bucket["unchanged"] += 1
continue
session.add(row.model(**row.values))
session.flush()
bucket["created"] += 1
return counts
def run_import(
source_root: Path,
database_url: str | None,
*,
apply: bool = False,
backup_dir: Path | None = None,
) -> dict[str, Any]:
"""Validate and optionally apply an import, returning metadata only."""
try:
snapshot = load_source(source_root)
rows = _build_rows(snapshot)
report: dict[str, Any] = {
"mode": "apply" if apply else "dry_run",
"source_checksum": _canonical_checksum(snapshot),
"counts": _counts(rows, apply=apply),
}
except MemoryError as exc:
raise ImportValidationError("source exceeds available memory") from exc
if not apply:
return report
if not database_url:
raise ImportValidationError("explicit database URL is required for apply")
if backup_dir is None:
raise ImportValidationError("backup directory is required for apply")
_backup_source(source_root.expanduser().resolve(), backup_dir)
try:
report["counts"] = _apply_rows(rows, database_url)
except SQLAlchemyError as exc:
raise ImportValidationError("target database operation failed") from exc
report["backup_created"] = True
return report
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Import JSON stores into PostgreSQL/SQLAlchemy schema")
parser.add_argument("--source", type=Path, required=True)
parser.add_argument("--database-url", default=os.environ.get("DATABASE_URL"))
parser.add_argument("--apply", action="store_true", help="write to the target database")
parser.add_argument("--backup-dir", type=Path)
args = parser.parse_args(argv)
try:
report = run_import(
args.source,
args.database_url,
apply=args.apply,
backup_dir=args.backup_dir,
)
except ImportValidationError:
print("import rejected: validation or target-conflict check failed", file=sys.stderr)
return 2
print(json.dumps(report, ensure_ascii=False, sort_keys=True))
return 0
if __name__ == "__main__":
raise SystemExit(main())