xinmi/scripts/weknora_category_etl.py

#!/usr/bin/env python3
"""Migrate legacy categories onto already imported acceptance documents.

Only metadata is changed. The SQLite source and the vector collections remain
read-only. A reviewed manifest is checked again inside the write transaction.
"""
from __future__ import annotations

import argparse
import datetime
import json
import os
import re
import subprocess
from pathlib import Path

from weknora_etl import plan, stable

DATABASE = "yawan_native_acceptance_20260915"
TENANT = 10059
PROVENANCE = "yawan_legacy_category_v1"


def literal(value):
    return "'" + str(value).replace("'", "''") + "'"


def sql_run(sql):
    result = subprocess.run([
        "docker", "--context", "colima-yawan", "exec", "-i", "WeKnora-postgres",
        "psql", "-X", "-qAt", "-v", "ON_ERROR_STOP=1", "-U", "postgres", "-d", DATABASE,
    ], input=sql, text=True, capture_output=True, timeout=60)
    if result.returncode:
        # SQL errors can contain values from the source; keep them off stdout.
        raise RuntimeError("Category transaction not confirmed; reconcile database and checkpoints before retrying")
    return result.stdout.strip()


def target_snapshot():
    return json.loads(sql_run(f"""
BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY;
SELECT json_build_object(
 'tenant', (SELECT json_build_object('id',id,'status',status) FROM tenants WHERE id={TENANT}),
 'knowledge_bases', (SELECT coalesce(json_agg(row_to_json(x) ORDER BY id),'[]') FROM
   (SELECT id,tenant_id,name FROM knowledge_bases WHERE tenant_id={TENANT} AND deleted_at IS NULL) x),
 'files', (SELECT coalesce(json_agg(row_to_json(x) ORDER BY id),'[]') FROM
   (SELECT id,tenant_id,knowledge_base_id,file_name,file_hash,file_size,folder_path,custom_metadata,parse_status
    FROM knowledges WHERE tenant_id={TENANT} AND deleted_at IS NULL) x),
 'tags', (SELECT coalesce(json_agg(row_to_json(x) ORDER BY id),'[]') FROM
   (SELECT id,tenant_id,knowledge_base_id,name,sort_order,deleted_at FROM knowledge_tags WHERE tenant_id={TENANT}) x),
 'relations', (SELECT coalesce(json_agg(row_to_json(x) ORDER BY knowledge_id,tag_id),'[]') FROM
   (SELECT r.* FROM knowledge_tag_relations r JOIN knowledges k ON k.id=r.knowledge_id
    WHERE k.tenant_id={TENANT}) x)
);
COMMIT;
"""))


def document_stem(name):
    path = Path(name)
    if path.name != name or path.suffix not in (".jsonl", ".md"):
        raise ValueError("Unexpected migration file name")
    return re.sub(r"\.part\d{3}$", "", path.stem)


def build_manifest(source, state, target):
    if source["source_mutated"] or not source["relation_validation"]["valid"]:
        raise ValueError("Source snapshot or relations are inconsistent")
    if state["tenant_id"] != TENANT or target["tenant"] != {"id": TENANT, "status": "active"}:
        raise ValueError("Wrong or inactive target tenant")
    source_kbs = {k["legacy_id"]: k for k in source["mapping"]["knowledge_bases"]}
    target_kbs = {k["id"]: k for k in target["knowledge_bases"]}
    if set(state["knowledge_bases"]) != {str(k) for k in source_kbs}:
        raise ValueError("Knowledge base mapping is incomplete")
    kb_map = {}
    approved_files = set()
    for legacy_id, kb in source_kbs.items():
        entry = state["knowledge_bases"][str(legacy_id)]
        native = target_kbs.get(entry["id"])
        if (not native or native["tenant_id"] != TENANT or native["name"] != kb["name"]
                or entry["legacy_name"] != kb["name"] or not entry["complete"] or entry["failed"]):
            raise ValueError("Knowledge base scope or completed import state differs")
        kb_map[kb["native_id"]] = entry["id"]
        for name in entry["uploaded"]:
            approved_files.add((entry["id"], str(Path(name).with_suffix(".md"))))
    if len(set(kb_map.values())) != len(kb_map):
        raise ValueError("Knowledge base mapping is not one-to-one")
    categories = {c["native_id"]: c for c in source["mapping"]["categories"]}

    def ancestry(category):
        result, seen = [], set()
        while category:
            cid = category["native_id"]
            if cid in seen:
                raise ValueError("Category cycle")
            seen.add(cid)
            name = category["name"]
            if not name.strip() or name != name.strip() or any(c in name for c in ("/", "\\", "\x00")) or name in (".", ".."):
                raise ValueError("Category name cannot be represented as a folder")
            result.insert(0, category)
            category = categories.get(category["parent_id"]) if category["parent_id"] else None
        return result

    tags = []
    tag_map = {}
    for cid, category in categories.items():
        chain = ancestry(category)
        path = "/".join(c["name"] for c in chain)
        tag_name = " / ".join(c["name"] for c in chain)
        if len(tag_name) > 128 or len(path) > 1024:
            raise ValueError("Category path exceeds target capacity")
        tag = {
            "id": stable(f"category:tenant:{TENANT}", category["legacy_id"]),
            "legacy_id": category["legacy_id"], "tenant_id": TENANT,
            "knowledge_base_id": kb_map[category["knowledge_base_id"]],
            "name": tag_name, "folder_path": path, "sort_order": category["sort_order"],
            "source": category,
        }
        tag_map[cid] = tag
        tags.append(tag)
    keys = [(t["knowledge_base_id"], t["name"]) for t in tags]
    if len(set(keys)) != len(keys):
        raise ValueError("Duplicate category paths")
    for old in target["tags"]:
        matches = [t for t in tags if t["id"] == old["id"] or
                   (t["knowledge_base_id"], t["name"]) == (old["knowledge_base_id"], old["name"])]
        for tag in matches:
            if old["deleted_at"] is not None or any(old[k] != tag[k] for k in
                    ("id", "tenant_id", "knowledge_base_id", "name", "sort_order")):
                raise ValueError("Existing tag conflicts with category import")

    grouped = {}
    actual_files = set()
    for file in target["files"]:
        if file["knowledge_base_id"] not in kb_map.values():
            continue
        if file["tenant_id"] != TENANT or file["parse_status"] != "completed":
            raise ValueError("Target file is outside the tenant or not indexed")
        if (file["knowledge_base_id"], file["file_name"]) not in approved_files:
            raise ValueError("Target file is absent from completed import state")
        actual_files.add((file["knowledge_base_id"], file["file_name"]))
        grouped.setdefault((file["knowledge_base_id"], document_stem(file["file_name"])), []).append(file)
    if actual_files != approved_files:
        raise ValueError("Completed import contains missing target files")
    files, relations, seen = [], [], set()
    for source_file in source["mapping"]["files"]:
        kb_id = kb_map[source_file["knowledge_base_id"]]
        candidates = grouped.get((kb_id, document_stem(source_file["file_name"])), [])
        if not candidates:
            raise ValueError("Legacy file has no imported document")
        if len(candidates) > 1:
            parts = sorted(int(m.group(1)) for f in candidates if
                           (m := re.search(r"\.part(\d{3})\.md$", f["file_name"])))
            if parts != list(range(1, len(candidates) + 1)):
                raise ValueError("Ambiguous or incomplete split file mapping")
        chain = ancestry(categories[source_file["category_id"]]) if source_file["category_id"] else []
        folder = "/".join(c["name"] for c in chain)
        marker = {"legacy_file_id": source_file["legacy_id"], "legacy_category_id": chain[-1]["legacy_id"] if chain else None,
                  "legacy_content_hash": source_file["content_hash"], "legacy_revision": source_file["revision"]}
        for target_file in candidates:
            if target_file["id"] in seen:
                raise ValueError("A document maps to multiple legacy files")
            seen.add(target_file["id"])
            metadata = target_file["custom_metadata"] or {}
            if target_file["folder_path"] not in ("", folder) or (PROVENANCE in metadata and metadata[PROVENANCE] != marker):
                raise ValueError("Existing document metadata conflicts with import")
            files.append({"id": target_file["id"], "knowledge_base_id": kb_id,
                          "file_name": target_file["file_name"], "legacy_id": source_file["legacy_id"],
                          "file_hash": target_file["file_hash"], "file_size": target_file["file_size"],
                          "folder_path": folder, "before_folder_path": target_file["folder_path"],
                          "before_metadata": metadata, "metadata": {**metadata, PROVENANCE: marker}})
            relations.extend({"knowledge_id": target_file["id"], "tag_id": tag_map[c["native_id"]]["id"]} for c in chain)
    if seen != {f["id"] for group in grouped.values() for f in group}:
        raise ValueError("Some imported files have no legacy mapping")
    return {"database": DATABASE, "tenant_id": TENANT, "source": source["source"],
            "knowledge_bases": list(kb_map.values()), "tags": tags, "files": files, "relations": relations,
            "summary": {"categories": len(tags), "source_files": len(source["mapping"]["files"]),
                        "target_files": len(files), "tag_relations": len(relations)}}


def transaction(manifest, commit=False):
    payload = literal(json.dumps(manifest, ensure_ascii=False))
    end = "COMMIT" if commit else "ROLLBACK"
    return f"""
BEGIN ISOLATION LEVEL SERIALIZABLE;
SET LOCAL statement_timeout = '30s';
SET LOCAL lock_timeout = '5s';
SET LOCAL standard_conforming_strings = on;
SELECT pg_advisory_xact_lock(10059, 20260916);
CREATE TEMP TABLE _category_import (data jsonb) ON COMMIT DROP;
INSERT INTO _category_import VALUES ({payload}::jsonb);
CREATE TEMP TABLE _tags ON COMMIT DROP AS SELECT t.* FROM _category_import m,
 jsonb_to_recordset(m.data->'tags') AS t(id text,legacy_id int,tenant_id int,knowledge_base_id text,name text,sort_order int,source jsonb);
CREATE TEMP TABLE _files ON COMMIT DROP AS SELECT f.* FROM _category_import m,
 jsonb_to_recordset(m.data->'files') AS f(id text,knowledge_base_id text,file_name text,file_hash text,file_size bigint,legacy_id int,folder_path text,before_folder_path text,before_metadata jsonb,metadata jsonb);
CREATE TEMP TABLE _relations ON COMMIT DROP AS SELECT r.* FROM _category_import m,
 jsonb_to_recordset(m.data->'relations') AS r(knowledge_id text,tag_id text);
DO $check$
BEGIN
 PERFORM 1 FROM tenants WHERE id={TENANT} FOR UPDATE;
 PERFORM 1 FROM knowledge_bases k WHERE k.id IN (SELECT knowledge_base_id FROM _tags) ORDER BY k.id FOR UPDATE;
 PERFORM 1 FROM knowledge_tags k JOIN _tags t ON k.id=t.id ORDER BY k.id FOR UPDATE OF k;
 IF current_database() <> '{DATABASE}' OR NOT EXISTS (SELECT 1 FROM tenants WHERE id={TENANT} AND status='active')
    OR EXISTS (SELECT 1 FROM _category_import WHERE (data->>'tenant_id')::int <> {TENANT} OR data->>'database' <> current_database())
 THEN RAISE EXCEPTION 'Wrong database or tenant'; END IF;
 IF EXISTS (SELECT 1 FROM _tags t LEFT JOIN knowledge_bases k ON k.id=t.knowledge_base_id
   WHERE t.tenant_id <> {TENANT} OR k.id IS NULL OR k.tenant_id <> {TENANT} OR k.deleted_at IS NOT NULL)
 THEN RAISE EXCEPTION 'Knowledge base scope changed'; END IF;
 PERFORM 1 FROM knowledges k JOIN _files f ON f.id=k.id FOR UPDATE OF k;
 IF EXISTS (SELECT 1 FROM _files f LEFT JOIN knowledges k ON k.id=f.id
   WHERE k.id IS NULL OR k.tenant_id <> {TENANT} OR k.knowledge_base_id <> f.knowledge_base_id
     OR k.file_name <> f.file_name OR k.deleted_at IS NOT NULL OR k.parse_status <> 'completed'
     OR k.file_hash IS DISTINCT FROM f.file_hash OR k.file_size IS DISTINCT FROM f.file_size
     OR k.folder_path IS DISTINCT FROM f.before_folder_path OR coalesce(k.custom_metadata,'{{}}') <> f.before_metadata)
 THEN RAISE EXCEPTION 'Document scope or metadata changed'; END IF;
 IF EXISTS (SELECT 1 FROM _tags t JOIN knowledge_tags k ON k.id=t.id OR
     (k.tenant_id=t.tenant_id AND k.knowledge_base_id=t.knowledge_base_id AND k.name=t.name)
   WHERE k.id <> t.id OR k.tenant_id <> t.tenant_id OR k.knowledge_base_id <> t.knowledge_base_id
     OR k.name <> t.name OR k.sort_order <> t.sort_order OR k.deleted_at IS NOT NULL)
 THEN RAISE EXCEPTION 'Category conflict'; END IF;
 IF EXISTS (SELECT 1 FROM _relations r LEFT JOIN _files f ON f.id=r.knowledge_id LEFT JOIN _tags t ON t.id=r.tag_id
   WHERE f.id IS NULL OR t.id IS NULL OR f.knowledge_base_id <> t.knowledge_base_id)
 THEN RAISE EXCEPTION 'Cross-knowledge-base relation'; END IF;
END $check$;
CREATE TABLE IF NOT EXISTS yawan_legacy_category_imports (
 tenant_id integer NOT NULL, legacy_id integer NOT NULL, native_id varchar(36) NOT NULL,
 source_data jsonb NOT NULL, source_snapshot_sha256 varchar(64) NOT NULL,
 created_at timestamptz NOT NULL DEFAULT now(), PRIMARY KEY(tenant_id,legacy_id)
);
DO $check$
BEGIN
 IF EXISTS (SELECT 1 FROM _tags t JOIN yawan_legacy_category_imports m ON m.tenant_id=t.tenant_id AND m.legacy_id=t.legacy_id
   WHERE m.native_id <> t.id OR m.source_data <> t.source)
 THEN RAISE EXCEPTION 'Legacy category provenance conflict'; END IF;
END $check$;
WITH inserted AS (
 INSERT INTO knowledge_tags (id,tenant_id,knowledge_base_id,name,sort_order)
 SELECT id,tenant_id,knowledge_base_id,name,sort_order FROM _tags
 ON CONFLICT (id) DO NOTHING RETURNING id
) SELECT json_build_object('inserted_tags',count(*)) FROM inserted;
WITH inserted AS (
 INSERT INTO knowledge_tag_relations (knowledge_id,tag_id) SELECT knowledge_id,tag_id FROM _relations
 ON CONFLICT DO NOTHING RETURNING knowledge_id
) SELECT json_build_object('inserted_relations',count(*)) FROM inserted;
WITH changed AS (
 UPDATE knowledges k SET folder_path=f.folder_path,custom_metadata=f.metadata,updated_at=now()
 FROM _files f WHERE k.id=f.id AND k.tenant_id={TENANT}
 AND (k.folder_path IS DISTINCT FROM f.folder_path OR k.custom_metadata IS DISTINCT FROM f.metadata)
 RETURNING k.id
) SELECT json_build_object('updated_documents',count(*)) FROM changed;
INSERT INTO yawan_legacy_category_imports (tenant_id,legacy_id,native_id,source_data,source_snapshot_sha256)
 SELECT t.tenant_id,t.legacy_id,t.id,t.source,m.data->'source'->>'sha256' FROM _tags t CROSS JOIN _category_import m
 ON CONFLICT DO NOTHING;
DO $check$
BEGIN
 IF EXISTS (SELECT 1 FROM _relations r LEFT JOIN knowledge_tag_relations k USING(knowledge_id,tag_id) WHERE k.knowledge_id IS NULL)
 THEN RAISE EXCEPTION 'Incomplete category relations'; END IF;
END $check$;
{end};
"""


def save_json(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    # Checkpoints are append-only so a retry cannot destroy the before-state.
    with path.open("x", encoding="utf-8") as stream:
        json.dump(value, stream, ensure_ascii=False, indent=2)
        stream.write("\n")
        stream.flush()
        os.fsync(stream.fileno())


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("plan", "apply"))
    parser.add_argument("--sqlite", type=Path, required=True)
    parser.add_argument("--state", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    if any(path.exists() for path in (args.out, *(args.out.with_suffix(suffix) for suffix in
            (".before.json", ".manifest.json", ".attempt.json", ".committed.json", ".after.json")))):
        parser.error("Output checkpoint already exists; use a new output name")
    source = plan(args.sqlite)
    state = json.loads(args.state.read_text(encoding="utf-8"))
    target = target_snapshot()
    manifest = build_manifest(source, state, target)
    save_json(args.out.with_suffix(".before.json"), target)
    save_json(args.out.with_suffix(".manifest.json"), manifest)
    dry_run = sql_run(transaction(manifest))
    result = {"checked_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
              "tenant_id": TENANT, "database": DATABASE, "summary": manifest["summary"],
              "source_snapshot_sha256": source["source"]["sha256"], "source_mutated": source["source_mutated"],
              "dry_run": [json.loads(line) for line in dry_run.splitlines() if line.startswith("{")],
              "committed": False}
    if args.command == "apply":
        save_json(args.out.with_suffix(".attempt.json"), {**result, "commit_requested": True,
                  "recovery": "Inspect category import ledger and rerun plan with a new output name if interrupted"})
        result["changes"] = [json.loads(line) for line in sql_run(transaction(manifest, commit=True)).splitlines() if line.startswith("{")]
        result["committed"] = True
        save_json(args.out.with_suffix(".committed.json"), result)
        after = target_snapshot()
        save_json(args.out.with_suffix(".after.json"), after)
        expected = {(r["knowledge_id"],r["tag_id"]) for r in manifest["relations"]}
        actual = {(r["knowledge_id"],r["tag_id"]) for r in after["relations"]}
        by_id = {f["id"]:f for f in after["files"]}
        result["verified"] = expected <= actual and all(by_id[f["id"]]["folder_path"] == f["folder_path"]
            and by_id[f["id"]]["custom_metadata"] == f["metadata"] for f in manifest["files"])
    save_json(args.out, result)
    print(json.dumps(result, ensure_ascii=False))
    return 0 if result.get("verified", True) else 1


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