MEGAINFRA-WEBAPP/scripts/pb-add-dn-staged-ingestion.py

#!/usr/bin/env python3
"""Add the staged delivery-note ingestion schema.

The migration is idempotent and is a dry run unless ``--apply`` is supplied.
It reads existing PocketBase fields and indexes before appending anything, so
running it against a live instance preserves unrelated schema changes.
"""
import re
import sys

from pb_env import PB, http, superuser_token


DRY_RUN = "--dry-run" in sys.argv or "--apply" not in sys.argv
APPLY = "--apply" in sys.argv


def confirmation_host():
    """Return the exact host confirmation supplied for an apply run."""
    for index, argument in enumerate(sys.argv):
        if argument.startswith("--confirm-host="):
            return argument.split("=", 1)[1]
        if argument == "--confirm-host" and index + 1 < len(sys.argv):
            return sys.argv[index + 1]
    return None


def get_collection(token, name):
    code, response = http("GET", f"/api/collections/{name}", token=token)
    return response if code == 200 else None


def require_collection(token, name):
    collection = get_collection(token, name)
    if collection is None:
        sys.exit(f"collection {name} not found")
    return collection


def relation(name, collection_id, required=False):
    return {
        "name": name,
        "type": "relation",
        "required": required,
        "collectionId": collection_id,
        "maxSelect": 1,
        "cascadeDelete": False,
    }


def field(name, field_type="text", **options):
    return {"name": name, "type": field_type, "required": False, **options}


def index_name(statement):
    """Extract an SQLite/PocketBase index name from a CREATE INDEX statement."""
    match = re.search(r"\bINDEX\s+(?:`([^`]+)`|\"([^\"]+)\"|([^\s(]+))", statement, re.IGNORECASE)
    if not match:
        raise ValueError(f"not a CREATE INDEX statement: {statement!r}")
    return next(value for value in match.groups() if value is not None).strip().lower()


def normalise_index_sql(statement):
    """Compare index SQL by meaning, not PocketBase's formatting choices."""
    value = statement.strip().rstrip(";")
    chunks = []
    literals = []
    index = 0
    while index < len(value):
        character = value[index]
        if character == "'":
            start = index
            index += 1
            while index < len(value):
                if value[index] == "'":
                    if index + 1 < len(value) and value[index + 1] == "'":
                        index += 2
                        continue
                    index += 1
                    break
                index += 1
            literals.append(value[start:index])
            chunks.append(f"\x00{len(literals) - 1}\x00")
            continue
        if character in '`"':
            quote = character
            index += 1
            identifier = []
            while index < len(value):
                if value[index] == quote:
                    if index + 1 < len(value) and value[index + 1] == quote:
                        identifier.append(quote)
                        index += 2
                        continue
                    index += 1
                    break
                identifier.append(value[index])
                index += 1
            chunks.append(f"`{''.join(identifier).lower()}`")
            continue
        if character.isalnum() or character in "_$":
            end = index + 1
            while end < len(value) and (value[end].isalnum() or value[end] in "_$"):
                end += 1
            chunks.append(value[index:end].lower())
            index = end
            continue
        chunks.append(character)
        index += 1

    value = "".join(chunks)
    value = re.sub(r"\s+", " ", value).strip().lower()
    # SQLite/PocketBase may insert or remove spaces around these tokens while
    # preserving the same columns, uniqueness and WHERE predicate.
    value = re.sub(r"\s*([(),=<>!])\s*", r"\1", value)
    # Quoted identifiers and safe unquoted identifiers share one canonical
    # representation; single-quoted literals are placeholders until restore.
    value = value.replace("`", "")
    for index, literal in enumerate(literals):
        value = value.replace(f"\x00{index}\x00", literal)
    return value


def merge_indexes(existing, wanted):
    """Union indexes by name, replacing stale same-name definitions in place."""
    merged = list(existing)
    positions = {index_name(statement): i for i, statement in enumerate(merged)}
    changes = []
    for statement in wanted:
        name = index_name(statement)
        position = positions.get(name)
        if position is None:
            positions[name] = len(merged)
            merged.append(statement)
            changes.append(("add", statement))
        elif normalise_index_sql(merged[position]) != normalise_index_sql(statement):
            merged[position] = statement
            changes.append(("replace", statement))
    return merged, changes


def add_fields(token, name, wanted):
    collection = require_collection(token, name)
    existing = list(collection.get("fields", []))
    names = {item.get("name") for item in existing}
    additions = [item for item in wanted if item["name"] not in names]
    if not additions:
        print(f"  {name}: fields already present")
        return
    print(f"  {name}: add fields {[item['name'] for item in additions]}")
    if not DRY_RUN:
        code, response = http(
            "PATCH", f"/api/collections/{collection['id']}",
            {"fields": existing + additions}, token,
        )
        if code != 200:
            sys.exit(f"PATCH {name} failed: {code} {response}")


def add_indexes(token, name, wanted):
    collection = require_collection(token, name)
    existing = list(collection.get("indexes", []) or [])
    merged, changes = merge_indexes(existing, wanted)
    if not changes:
        print(f"  {name}: indexes already present")
        return
    print(f"  {name}: reconcile indexes {[item for _, item in changes]}")
    if not DRY_RUN:
        code, response = http(
            "PATCH", f"/api/collections/{collection['id']}",
            {"indexes": merged}, token,
        )
        if code != 200:
            sys.exit(f"PATCH {name} indexes failed: {code} {response}")


def create_collection(token, name, fields, indexes):
    existing = get_collection(token, name)
    if existing is not None:
        print(f"  {name}: already exists")
        add_fields(token, name, fields)
        add_indexes(token, name, indexes)
        return existing["id"]
    print(f"  {name}: create collection")
    if DRY_RUN:
        return name
    code, response = http(
        "POST", "/api/collections",
        {"name": name, "type": "base", "fields": fields, "indexes": indexes},
        token,
    )
    if code != 200:
        sys.exit(f"CREATE {name} failed: {code} {response}")
    return response["id"]


def system_fields():
    return [
        field("created", "date", onCreate=True),
        field("updated", "date", onCreate=True, onUpdate=True),
    ]


def supplier_identity_fields():
    """Optional supplier identifiers used by document entity resolution."""
    return [field("registration_number"), field("vat_number")]


def self_test():
    identity_fields = supplier_identity_fields()
    assert [item["name"] for item in identity_fields] == [
        "registration_number", "vat_number",
    ]
    assert all(item["required"] is False for item in identity_fields)

    old = [
        "CREATE INDEX `unrelated` ON `x` (`value`)",
        "CREATE UNIQUE INDEX `idx_delivery_notes_batch_note` ON `delivery_notes` (`batch_note`)",
    ]
    replacement = (
        "CREATE UNIQUE INDEX `idx_delivery_notes_batch_note` ON `delivery_notes` (`batch_note`) "
        "WHERE `batch_note` IS NOT NULL AND `batch_note` != ''"
    )
    merged, changes = merge_indexes(old, [replacement])
    assert merged == [old[0], replacement]
    assert changes == [("replace", replacement)]
    assert merge_indexes(merged, [replacement])[1] == []
    candidate_key = field("candidate_key")
    assert candidate_key == {"name": "candidate_key", "type": "text", "required": False}
    global_page = field("global_page", "number")
    assert global_page == {"name": "global_page", "type": "number", "required": False}
    candidate_index = (
        "CREATE UNIQUE INDEX `idx_dn_batch_notes_candidate_key` "
        "ON `dn_batch_notes` (`batch`, `candidate_key`) "
        "WHERE `candidate_key` IS NOT NULL AND `candidate_key` != ''"
    )
    candidate_indexes, candidate_changes = merge_indexes([], [candidate_index])
    assert candidate_indexes == [candidate_index]
    assert candidate_changes == [("add", candidate_index)]
    assert merge_indexes(candidate_indexes, [candidate_index])[1] == []
    # PocketBase returns the same three partial indexes uppercased and with
    # multiline/variable whitespace. Those rows must be idempotent.
    semantic_wanted = [
        "CREATE UNIQUE INDEX `idx_dn_batch_notes_candidate_key` ON `dn_batch_notes` (`batch`, `candidate_key`) WHERE `candidate_key` IS NOT NULL AND `candidate_key` != ''",
        "CREATE UNIQUE INDEX `idx_delivery_notes_batch_note` ON `delivery_notes` (`batch_note`) WHERE `batch_note` IS NOT NULL AND `batch_note` != ''",
        "CREATE UNIQUE INDEX `idx_dn_batches_capture_session_id` ON `dn_batches` (`capture_session_id`) WHERE `capture_session_id` != ''",
    ]
    pocketbase_stored = [
        "CREATE UNIQUE INDEX `IDX_DN_BATCH_NOTES_CANDIDATE_KEY`\nON `DN_BATCH_NOTES` (`BATCH`,  `CANDIDATE_KEY`)\nWHERE `CANDIDATE_KEY` IS NOT NULL AND `CANDIDATE_KEY` != ''",
        "CREATE UNIQUE INDEX `IDX_DELIVERY_NOTES_BATCH_NOTE` ON `DELIVERY_NOTES`\n( `BATCH_NOTE` ) WHERE `BATCH_NOTE` IS NOT NULL\nAND `BATCH_NOTE` != ''",
        "CREATE UNIQUE INDEX `IDX_DN_BATCHES_CAPTURE_SESSION_ID` ON `DN_BATCHES` (`CAPTURE_SESSION_ID`)\nWHERE `CAPTURE_SESSION_ID` != ''",
    ]
    assert merge_indexes(pocketbase_stored, semantic_wanted)[1] == []
    near_misses = [
        semantic_wanted[0].replace("candidate_key`)", "other_key`)") ,
        semantic_wanted[1].replace("CREATE UNIQUE", "CREATE"),
        semantic_wanted[2].replace(" != ''", " IS NOT NULL"),
    ]
    for stored, near_miss in zip(pocketbase_stored, near_misses):
        assert merge_indexes([stored], [near_miss])[1] == [("replace", near_miss)]
    identifier_styles = (
        "CREATE UNIQUE INDEX `idx_style` ON `Delivery_Notes` (`Batch_Note`) WHERE `Batch_Note` != ''",
        'create unique index "IDX_STYLE" on "DELIVERY_NOTES" ("BATCH_NOTE") where "BATCH_NOTE" != \'\'',
        "CREATE UNIQUE INDEX idx_style ON delivery_notes (batch_note) WHERE batch_note != ''",
    )
    assert merge_indexes([identifier_styles[0]], [identifier_styles[1]])[1] == []
    assert merge_indexes([identifier_styles[0]], [identifier_styles[2]])[1] == []
    literal_near_miss = identifier_styles[2].replace("!= ''", "!= ' '")
    assert merge_indexes([identifier_styles[0]], [literal_near_miss])[1] == [("replace", literal_near_miss)]
    print("self-test: index replacement and idempotence passed")
    return 0


def main():
    if "--self-test" in sys.argv:
        return self_test()
    if APPLY and confirmation_host() != PB:
        sys.exit(
            "REFUSING --apply: pass --confirm-host with the exact configured "
            f"PocketBase host ({PB})"
        )
    try:
        token = superuser_token()
    except SystemExit as error:
        # A dry run is useful in CI and during review even when no PB
        # credentials are mounted.  It still reports the complete intended
        # mutation and, by definition, performs no writes.
        if not DRY_RUN:
            raise
        print(f"offline preview: {error}")
        print("would create: dn_batch_sources, dn_batch_pages, "
              "supplier_document_aliases, dn_review_events, dn_pipeline_locks")
        print("would extend: dn_batches, dn_batch_notes(candidate_key), delivery_notes")
        print("would add unique indexes: source(batch,file_hash), "
              "page(source,source_page,region) and page(batch,global_page), alias(normalised_key), "
              "lock(batch), candidate(batch,candidate_key), delivery_notes(batch_note)")
        print("would add partial unique capture-session index on dn_batches")
        print("\nDRY RUN. No schema changes written.")
        return 0
    print(f"target: {PB}   mode: {'DRY RUN' if DRY_RUN else 'APPLY'}")

    batches = require_collection(token, "dn_batches")
    notes = require_collection(token, "dn_batch_notes")
    suppliers = require_collection(token, "suppliers")
    delivery_notes = require_collection(token, "delivery_notes")
    batch_id = batches["id"]
    notes_id = notes["id"]
    supplier_id = suppliers["id"]
    delivery_notes_id = delivery_notes["id"]

    # Supplier identity evidence is optional: existing rows remain valid and
    # operators can backfill identifiers as documents make them available.
    add_fields(token, "suppliers", supplier_identity_fields())

    source_id = create_collection(
        token, "dn_batch_sources",
        system_fields() + [
            relation("batch", batch_id, required=True),
            field("file", "file", required=True, maxSelect=1, maxSize=20971520),
            field("file_hash", required=True),
            # Not required: PocketBase treats 0 as blank for a required
            # number, and the first source in a batch is legitimately 0.
            field("ordinal", "number"),
            field("mime", required=True),
            field("original_name", required=True), field("capture_session_id"),
            field("delivery_session_id"), field("captured_at", "date"),
            field("latitude", "number"), field("longitude", "number"),
            field("location_accuracy_m", "number"),
        ],
        ["CREATE UNIQUE INDEX `idx_dn_batch_sources_batch_file_hash` "
         "ON `dn_batch_sources` (`batch`, `file_hash`)"]
    )
    create_collection(
        token, "dn_batch_pages",
        system_fields() + [
            relation("batch", batch_id, required=True),
            relation("source", source_id, required=True),
            field("source_page", "number", required=True),
            # Assigned by the staged runner after normalisation. Optional so
            # legacy rows can be backfilled only when their identity is clear.
            field("global_page", "number"),
            # Not required for the same reason as `ordinal`: a whole-page
            # read carries region 0, which a required number would reject.
            field("region", "number"),
            field("image", "file", required=True, maxSelect=1, maxSize=20971520),
            field("image_hash", required=True), field("mime", required=True),
            field("width", "number"),
            field("height", "number"), field("rotation", "number"),
            field("read_status"), field("extraction", "json", maxSize=2000000),
            field("schema_version"), field("model"), field("prompt_version"),
            field("read_error"), field("manual_group"),
        ],
        ["CREATE UNIQUE INDEX `idx_dn_batch_pages_source_page_region` "
         "ON `dn_batch_pages` (`source`, `source_page`, `region`)",
         "CREATE UNIQUE INDEX `idx_dn_batch_pages_batch_global_page` "
         "ON `dn_batch_pages` (`batch`, `global_page`) "
         "WHERE `global_page` IS NOT NULL"]
    )
    create_collection(
        token, "supplier_document_aliases",
        system_fields() + [
            relation("supplier", supplier_id, required=True),
            field("raw_value", required=True),
            field("normalised_key", required=True), field("source"), field("approved_by"),
        ],
        ["CREATE UNIQUE INDEX `idx_supplier_document_aliases_normalised_key` "
         "ON `supplier_document_aliases` (`normalised_key`)"]
    )
    create_collection(
        token, "dn_review_events",
        system_fields() + [
            relation("batch", batch_id, required=True), relation("note", notes_id),
            field("event_key"),
            field("action"), field("before", "json", maxSize=2000000),
            field("after", "json", maxSize=2000000), field("reason"), field("actor"),
        ], ["CREATE UNIQUE INDEX `idx_dn_review_events_event_key` "
            "ON `dn_review_events` (`event_key`) "
            "WHERE `event_key` IS NOT NULL AND `event_key` != ''"],
    )
    create_collection(
        token, "dn_pipeline_locks",
        system_fields() + [relation("batch", batch_id, required=True),
                           field("expires_at", "date", required=True),
                           field("worker", required=True), field("scope_key", max=100)],
        ["CREATE UNIQUE INDEX `idx_dn_pipeline_locks_batch` "
         "ON `dn_pipeline_locks` (`batch`)",
         "CREATE UNIQUE INDEX idx_dn_posting_scope ON dn_pipeline_locks (scope_key) WHERE scope_key != ''"]
    )

    add_fields(token, "dn_batches", [
        field("stage", "select", maxSelect=1,
              values=["queued", "preparing", "reading", "split_review", "complete", "failed"]),
        field("source_count", "number"), field("processing_completed_at", "date"),
        field("capture_session_id"), field("assembly_revision", "number"),
    ])
    add_fields(token, "dn_batch_notes", [
        field("stage", "select", maxSelect=1,
              values=["matching", "needs_review", "ready_to_post", "posted", "dismissed"]),
        field("assembly_revision", "number"), field("pages", "json", maxSize=2000000),
        relation("supplier", supplier_id), field("supplier_name_raw"),
        field("match_method"), field("match_evidence", "json", maxSize=2000000), field("candidate_key"),
        field("posted_at", "date"),
        # What a person set by hand, kept apart from what the model read.
        # Every run rebuilds a candidate from its page reads, so without
        # this an operator's corrected order number is overwritten by the
        # original misread on the very next run.
        field("manual_overrides", "json", maxSize=100000),
        field("reconciliation_preview", "json", maxSize=2000000),
        field("line_posting", "json", maxSize=2000000),
    ])
    add_indexes(token, "dn_batch_notes", [
        "CREATE UNIQUE INDEX `idx_dn_batch_notes_candidate_key` "
        "ON `dn_batch_notes` (`batch`, `candidate_key`) "
        "WHERE `candidate_key` IS NOT NULL AND `candidate_key` != ''",
    ])
    add_fields(token, "delivery_notes", [relation("batch_note", notes_id)])

    add_indexes(token, "delivery_notes", [
        "CREATE UNIQUE INDEX `idx_delivery_notes_batch_note` "
        "ON `delivery_notes` (`batch_note`) "
        "WHERE `batch_note` IS NOT NULL AND `batch_note` != ''",
    ])
    add_indexes(token, "dn_batches", [
        "CREATE UNIQUE INDEX `idx_dn_batches_capture_session_id` "
        "ON `dn_batches` (`capture_session_id`) WHERE `capture_session_id` != ''",
    ])

    print("\nDRY RUN. No schema changes written." if DRY_RUN else "\nSchema applied.")
    return 0


if __name__ == "__main__":
    sys.exit(main())