"""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())