MEGAINFRA-WEBAPP/scripts/migrations/2026-08-19-extend-data-import-batches.py

#!/usr/bin/env python3
"""
One-shot PocketBase migration applied via:

  PB_URL=https://pb.mega.infra.apphub.co.za \\
  PB_ADMIN_EMAIL=admin@mega.co.za \\
  PB_ADMIN_PASSWORD=... \\
  python3 scripts/migrations/2026-08-19-extend-data-import-batches.py

Extends the data_import_batches collection with the async-queue fields
required by the redesigned PO/DN import flow (see spec
2026-08-19-po-import-queue-redesign.md).

Idempotent: re-running is safe; existing fields are skipped.

NOTE: This script deviates from the plan in two places because PB's
v0.39 API uses `fields` (not `schema`) as the key, and field options
(`max`, `min`, `values`, `maxSelect`) live at the field top level rather
than nested under `options`. Both shapes are reflected here.
"""
import json
import os
import sys
import urllib.error
import urllib.request


COLLECTION = "data_import_batches"

# Fields to add. PB v0.39 shape: `options`-equivalent keys (`max`, `min`,
# `values`, `maxSelect`, `noDecimal`) live at the field top level.
NEW_FIELDS = {
    "content_type": {
        "name": "content_type",
        "type": "text",
        "required": False,
        "max": 64,
    },
    "status": {
        "name": "status",
        "type": "select",
        "required": False,  # legacy rows have no value
        "maxSelect": 1,
        "values": ["pending", "processing", "processed", "failed", "dead"],
    },
    "attempts": {
        "name": "attempts",
        "type": "number",
        "required": False,
        "min": 0,
        "max": 100,
        "noDecimal": True,
    },
    "last_error": {
        "name": "last_error",
        "type": "text",
        "required": False,
        "max": 2000,
    },
    "last_attempt_at": {
        "name": "last_attempt_at",
        "type": "date",
        "required": False,
    },
    "next_retry_at": {
        "name": "next_retry_at",
        "type": "date",
        "required": False,
    },
    "processed_at": {
        "name": "processed_at",
        "type": "date",
        "required": False,
    },
    "processed_record_id": {
        "name": "processed_record_id",
        "type": "text",
        "required": False,
        "max": 32,
    },
}

# Indexes to ensure exist. PB appends these into the collection's
# `indexes` array; the SQL is executed when the collection is saved.
# PB stores indices with the table name in upper-case and reformats the
# SQL string, so the idempotency check uses the index NAME (`idx_*`)
# instead of comparing the full SQL.
NEW_INDEX_NAME = "idx_batches_pending"
NEW_INDEX = (
    f"CREATE INDEX `{NEW_INDEX_NAME}` ON `data_import_batches` "
    "(`status`, `batch_type`, `next_retry_at`)"
)


def admin_auth(base_url: str) -> str:
    email = os.environ["PB_ADMIN_EMAIL"]
    password = os.environ["PB_ADMIN_PASSWORD"]
    body = json.dumps({"identity": email, "password": password}).encode()
    req = urllib.request.Request(
        f"{base_url}/api/collections/_superusers/auth-with-password",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req) as resp:
        return json.loads(resp.read())["token"]


def get_collection(base_url: str, token: str, name: str):
    req = urllib.request.Request(
        f"{base_url}/api/collections/{name}",
        headers={"Authorization": token},
    )
    try:
        with urllib.request.urlopen(req) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        if e.code == 404:
            return None
        raise


def patch_collection(base_url: str, token: str, name: str, body: dict):
    req = urllib.request.Request(
        f"{base_url}/api/collections/{name}",
        data=json.dumps(body).encode(),
        headers={
            "Authorization": token,
            "Content-Type": "application/json",
        },
        method="PATCH",
    )
    try:
        with urllib.request.urlopen(req) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        body_txt = e.read().decode() if e.fp else ""
        raise RuntimeError(f"PATCH {name} failed: {e.code} {body_txt}") from e


def main() -> int:
    base_url = os.environ.get("PB_URL", "http://127.0.0.1:8090").rstrip("/")
    token = admin_auth(base_url)
    coll = get_collection(base_url, token, COLLECTION)
    if not coll:
        print(f"ERROR: collection {COLLECTION!r} not found on {base_url}")
        return 1

    existing_fields = coll.get("fields", [])
    existing_field_names = {f["name"] for f in existing_fields}
    new_fields = [
        NEW_FIELDS[name] for name in NEW_FIELDS
        if name not in existing_field_names
    ]
    if new_fields:
        merged_fields = existing_fields + new_fields
        print(f"Adding {len(new_fields)} new field(s) to {COLLECTION}: "
              f"{[f['name'] for f in new_fields]}")
        patch_collection(base_url, token, COLLECTION, {"fields": merged_fields})
        # Re-read so we have the updated state for subsequent steps.
        coll = get_collection(base_url, token, COLLECTION)
        existing_fields = coll.get("fields", [])
    else:
        print(f"No-op: all NEW_FIELDS already present on {COLLECTION}.")

    # Extend batch_type enum if needed.
    schema_bt = next((f for f in existing_fields if f["name"] == "batch_type"), None)
    if schema_bt:
        bt_values = set(schema_bt.get("values", []))
        wanted = {"purchase_orders", "delivery_notes"}
        missing = wanted - bt_values
        if missing:
            # Build the merged fields array and PATCH with the new values.
            merged_fields = []
            for f in existing_fields:
                if f["name"] == "batch_type":
                    f = dict(f)
                    # Sort for deterministic order; keep 'timesheets' first
                    # to preserve the legacy list ordering visually.
                    new_order = list(bt_values | wanted)
                    f["values"] = new_order
                merged_fields.append(f)
            print(f"Extending batch_type enum with {sorted(missing)}")
            patch_collection(
                base_url, token, COLLECTION, {"fields": merged_fields}
            )
            coll = get_collection(base_url, token, COLLECTION)
            existing_fields = coll.get("fields", [])
        else:
            print("No-op: batch_type enum already contains purchase_orders "
                  "and delivery_notes.")

    # Ensure the index exists. The PB collection's `indexes` is a list of
    # raw SQL strings that PocketBase executes when the collection is
    # saved. PB reformats the SQL on storage (uppercases the table name
    # and reformats whitespace), so we check by index NAME rather than
    # by exact SQL string equality.
    existing_indexes = coll.get("indexes", []) or []
    if any(NEW_INDEX_NAME in idx for idx in existing_indexes):
        print(f"No-op: index {NEW_INDEX_NAME!r} already present.")
    else:
        print(f"Adding index: {NEW_INDEX}")
        patch_collection(
            base_url, token, COLLECTION,
            {"indexes": list(existing_indexes) + [NEW_INDEX]},
        )

    # Final sanity dump for the operator.
    final = get_collection(base_url, token, COLLECTION)
    final_field_names = sorted(f["name"] for f in final.get("fields", []))
    print(f"\nFinal field list on {COLLECTION}: {final_field_names}")
    bt = next((f for f in final["fields"] if f["name"] == "batch_type"), None)
    if bt:
        print(f"batch_type.values = {bt.get('values')}")
    print(f"indexes = {final.get('indexes', [])}")
    print("\nMigration complete.")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyError as e:
        print(f"ERROR: missing environment variable {e}. "
              "Set PB_ADMIN_EMAIL and PB_ADMIN_PASSWORD.", file=sys.stderr)
        sys.exit(2)