xinmi/scripts/weknora_etl.py

#!/usr/bin/env python3
"""Read-only legacy migration planner.

The source is never opened with SQLite's ``immutable=1`` flag: that flag skips
WAL pages and can return a database state that never existed.  Commands first
make a coherent temporary copy using SQLite's online backup API, then read
only that copy.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import sqlite3
import tempfile
import uuid
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator

NAMESPACE = uuid.UUID("7e8c2d45-1a67-4b2f-9d94-9dbb7e8e5f91")
TABLES = ("appuser", "knowledgebase", "knowledgefile", "knowledgecategory", "agent", "agentskill", "agentskillbinding", "canvas", "trainingmap", "segment", "digitalhuman", "memory", "knowledgeusage", "pointaccount", "pointtransaction", "auditlog", "systemlog")


def open_readonly(path: Path) -> sqlite3.Connection:
    """Open a live database read-only while allowing SQLite to read its WAL."""
    # as_uri escapes spaces and '#' in mounted volume paths correctly. mode=ro
    # is required: query_only alone still allows SQLite to create sidecars.
    db = sqlite3.connect(f"{path.resolve().as_uri()}?mode=ro", uri=True)
    db.execute("PRAGMA query_only=ON")
    return db


@contextmanager
def consistent_snapshot(path: Path) -> Iterator[tuple[sqlite3.Connection, dict]]:
    """Yield a consistent read-only snapshot and remove it on exit."""
    before = path.stat()
    source = open_readonly(path)
    try:
        with tempfile.TemporaryDirectory(prefix="weknora-etl-") as tmp:
            snapshot_path = Path(tmp) / "snapshot.sqlite"
            target = sqlite3.connect(snapshot_path)
            try:
                # backup() includes committed pages in both the main file and
                # its WAL. It does not modify the source database.
                source.backup(target, pages=256, sleep=0.05)
                target.commit()
            finally:
                # Keep this connection open for the read phase.  It was
                # created in the temporary directory and has no source side
                # effects; query_only prevents accidental writes.
                target.execute("PRAGMA query_only=ON")
            try:
                digest = hashlib.sha256()
                with snapshot_path.open("rb") as stream:
                    for chunk in iter(lambda: stream.read(1024 * 1024), b""):
                        digest.update(chunk)
                info = {"snapshot_method": "sqlite_backup", "source_bytes": before.st_size, "source_mtime_ns": before.st_mtime_ns, "bytes": snapshot_path.stat().st_size, "sha256": digest.hexdigest(), "source_mutated": path.stat().st_mtime_ns != before.st_mtime_ns}
                yield target, info
            finally:
                target.close()
    finally:
        source.close()


def rows(db: sqlite3.Connection, table: str, columns: str):
    return db.execute(f"SELECT {columns} FROM {table} ORDER BY id").fetchall()


def stable(kind: str, source_id: int) -> str:
    return str(uuid.uuid5(NAMESPACE, f"legacy:{kind}:{source_id}"))


def _audit_snapshot(db: sqlite3.Connection, db_path: Path, snapshot_info: dict) -> dict:
    integrity = db.execute("PRAGMA integrity_check").fetchone()[0]
    if integrity != "ok":
        raise RuntimeError(f"SQLite integrity check failed: {integrity}")
    counts: dict[str, int] = {}; missing: list[str] = []
    for table in TABLES:
        try:
            counts[table] = db.execute(f"SELECT count(*) FROM {table}").fetchone()[0]
        except sqlite3.OperationalError as exc:
            if "no such table" in str(exc).lower(): counts[table] = 0; missing.append(table)
            else: raise
    return {"path": str(db_path), **snapshot_info, "integrity": integrity, "counts": counts, "missing_tables": missing}


def audit(db_path: Path) -> dict:
    with consistent_snapshot(db_path) as (db, info): return _audit_snapshot(db, db_path, info)


def _safe_rows(db: sqlite3.Connection, table: str, columns: str):
    try: return rows(db, table, columns)
    except sqlite3.OperationalError as exc:
        if "no such table" in str(exc).lower(): return []
        raise


def _validate_relations(kbs, categories, files) -> dict:
    kb_ids = {r[0] for r in kbs}; category_by_id = {r[0]: r for r in categories}; category_errors: list[dict] = []
    for row in categories:
        cid, kb_id, parent_id = row[:3]
        if kb_id not in kb_ids: category_errors.append({"type": "missing_knowledge_base", "category_id": cid, "kb_id": kb_id})
        if parent_id is not None and parent_id not in category_by_id: category_errors.append({"type": "missing_parent", "category_id": cid, "parent_id": parent_id})
        elif parent_id is not None and category_by_id[parent_id][1] != kb_id: category_errors.append({"type": "cross_knowledge_base_parent", "category_id": cid, "parent_id": parent_id})
        seen: set[int] = set(); cursor = cid
        while cursor in category_by_id:
            if cursor in seen: category_errors.append({"type": "category_cycle", "category_id": cid}); break
            seen.add(cursor); cursor = category_by_id[cursor][2]
            if cursor is None: break
    file_errors: list[dict] = []
    for row in files:
        fid, kb_id, _, category_id = row[:4]
        if kb_id not in kb_ids: file_errors.append({"type": "missing_knowledge_base", "file_id": fid, "kb_id": kb_id})
        if category_id is not None:
            category = category_by_id.get(category_id)
            if category is None: file_errors.append({"type": "missing_category", "file_id": fid, "category_id": category_id})
            elif category[1] != kb_id: file_errors.append({"type": "cross_knowledge_base_category", "file_id": fid, "category_id": category_id})
    return {"categories": category_errors, "files": file_errors, "valid": not category_errors and not file_errors}


def plan(db_path: Path) -> dict:
    with consistent_snapshot(db_path) as (db, info):
        kbs = _safe_rows(db, "knowledgebase", "id, name, desc, organization_id, owner_user_id, visibility, source_dir")
        categories = _safe_rows(db, "knowledgecategory", "id, kb_id, parent_id, name, desc, sort_order")
        files = _safe_rows(db, "knowledgefile", "id, kb_id, file_name, category_id, content_hash, revision")
        result = {"source": _audit_snapshot(db, db_path, info), "mode": "dry-run", "writes_performed": False, "source_mutated": info["source_mutated"], "relation_validation": _validate_relations(kbs, categories, files), "mapping": {"knowledge_bases": [{"legacy_id": r[0], "native_id": stable("knowledge-base", r[0]), "name": r[1], "description": r[2], "organization_id": r[3], "owner_user_id": r[4], "visibility": r[5], "source_dir": r[6]} for r in kbs], "categories": [{"legacy_id": r[0], "native_id": stable("category", r[0]), "knowledge_base_id": stable("knowledge-base", r[1]), "parent_id": stable("category", r[2]) if r[2] is not None else None, "name": r[3], "description": r[4], "sort_order": r[5]} for r in categories], "files": [{"legacy_id": r[0], "native_id": stable("knowledge", r[0]), "knowledge_base_id": stable("knowledge-base", r[1]), "file_name": r[2], "category_id": stable("category", r[3]) if r[3] is not None else None, "content_hash": r[4], "revision": r[5], "action": "upload-and-reindex"} for r in files]}, "excluded_until_native_mapping": ["agents", "skills", "canvases", "training_maps", "segments", "billing", "logs"]}
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description="Create a read-only WeKnora legacy migration manifest")
    parser.add_argument("command", choices=("audit", "plan")); parser.add_argument("--sqlite", required=True, type=Path); parser.add_argument("--out", required=True, type=Path)
    args = parser.parse_args()
    if not args.sqlite.is_file(): parser.error(f"SQLite file not found: {args.sqlite}")
    payload = audit(args.sqlite) if args.command == "audit" else plan(args.sqlite)
    args.out.parent.mkdir(parents=True, exist_ok=True); args.out.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({"command": args.command, "out": str(args.out), "writes_performed": False}, ensure_ascii=False)); return 0


if __name__ == "__main__": raise SystemExit(main())