scripts.purge_memgraph_test_schema

Reference › scripts › purge_memgraph_test_schema

scripts/purge_memgraph_test_schema.py purges leaked per-test chunk-label DDL from a Memgraph instance: tests that create a label of their own leave constraints and indexes behind, and this script finds and drops them. A human runs it. It reads the canonical label set from orchestrata_ultimate.knowledge_graph.canonical_schema (_canonical_labels) so that production labels are never touched.

is_leaked_label is true for a per-test label and false for production and per-worker ones. survey maps each leaked label to the number of nodes still carrying it and counts the DDL, vector_indexes_by_label finds their vector indexes, constrained_properties and indexed_properties list what must be dropped for a label, and drop_label drops it. main takes --apply (default is a dry-run report), --uri (default the first of the configured URI keys) and --all-databases (sweep every database on the instance, since the leak spreads across per-worker databases). It prints "clean" when nothing is leaked; otherwise it says dry run, or reports the labels remaining after dropping. The exit code is 1 only when labels remain.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/scripts/purge_memgraph_test_schema.py

#!/usr/bin/env python3
"""Purge leaked per-test chunk-label DDL from a Memgraph instance.

``MemgraphSemanticGraph.connect()`` creates a uniqueness constraint and two
lookup indexes for whatever ``chunk_label`` its config names.
``clear()``/``clear_prefixed()`` delete NODES only — the schema survives, and
Memgraph replays every entry at boot recovery. When a test fixture builds its
label from the ``test_id`` uuid, that is unbounded growth.

It has bitten twice:

* 2026-07-15 — 3,379 leaked ``*__TEST_*`` labels (6,713 indexes + 3,323
  constraints guarding 9 nodes) cost the shared cluster ~17 minutes of pegged
  CPU per data-node restart. Purged by hand.
* 2026-07-26 — regrown to 824 labels (823 constraints + 277 indexes guarding
  1 node) because two architecture-twin fixtures never dropped their schema.

The hand-purge is why it regrew unnoticed: there was nothing to re-run and
nothing to measure. This script is the repeatable counterpart. The leak's
SOURCE is closed separately — the fixtures now call
``drop_chunk_label_schema()``, and ``scripts/test-review.py`` rule T17 fails
any future fixture that builds a uuid-scoped ``chunk_label`` without one.

Matching is deliberately narrow: a label qualifies only if it carries the
``__TEST_<hex>`` suffix that ``test_id``-derived labels produce. Production
labels (``SemanticChunk``, ``CodestralChunk``, ``Entity``, ``Page``, ...) and
the bounded per-worker labels (``TestChunkgw0``) can never match, so this
cannot touch real schema.

Dry-run by default; ``--apply`` drops. Idempotent — a clean instance reports
0 everywhere.
"""

from __future__ import annotations

import argparse
import os
import re
import sys
from collections.abc import Mapping
from functools import reduce
from typing import Any

# Every shape a per-test label takes, anchored at the END so a production
# label that merely CONTAINS a marker mid-string is never swept.
#
# The first pattern was the only one until 2026-08-06, when surveying all 31
# databases on the test cluster turned up 57 leaked constraints this narrow
# match could not see: five call sites build the label as
# ``f"{base}{test_id}"`` with a raw ``CREATE CONSTRAINT``, producing
# ``CanonTest_test_<hex>`` / ``CliApplyProbe_test_<hex>`` /
# ``EnsureSchema_test_<hex>`` / ``MigTestTEST<hex>`` rather than the adapter's
# normalized ``__TEST_<hex>``.
_TEST_LABEL_PATTERNS = (
    re.compile(r"__TEST_[0-9a-f]+$"),
    re.compile(r"_test_[0-9a-f]{8,}$"),
    re.compile(r"TEST[0-9a-f]{8,}$"),
)

_URI_KEYS = ("MEMGRAPH_URI", "TEST_MEMGRAPH_URI")


def _canonical_labels() -> frozenset[str]:
    """The declared canonical label set, or empty if unavailable.

    A second, independent guard: matching is already anchored to a per-test
    hex suffix that no canonical label carries, so this can only ever be
    belt-and-braces — but the cost of a wrong drop here is production schema,
    so the check is worth its two lines.
    """
    try:
        from orchestrata_ultimate.knowledge_graph.canonical_schema import (
            CANONICAL_LABELS,
        )
    except Exception:  # noqa: BLE001 - standalone use against a bare instance
        return frozenset()
    return frozenset(CANONICAL_LABELS)


def is_leaked_label(label: str) -> bool:
    """True for a per-test label, False for production and per-worker ones."""
    bare = label.lstrip(":")
    if bare in _canonical_labels():
        return False
    return any(pattern.search(bare) for pattern in _TEST_LABEL_PATTERNS)


def resolve_uri(env: Mapping[str, str], explicit: str | None) -> str:
    """Target URI: ``--uri`` wins, else the first populated env key.

    No literal address here — the addresses live in ``.env`` (RUL-003), and
    which instance to purge is an explicit choice the caller makes.
    """
    if explicit:
        return explicit
    for key in _URI_KEYS:
        value = env.get(key, "").strip()
        if value:
            return value
    raise SystemExit(
        f"no target: pass --uri or set one of {', '.join(_URI_KEYS)} in the environment"
    )


def survey(session: Any) -> tuple[dict[str, int], int]:
    """Leaked label -> number of nodes still carrying it, plus total DDL count."""

    def leaked(statement: str) -> list[dict[str, Any]]:
        rows = map(dict, session.run(statement))
        return [
            row for row in rows if is_leaked_label((row.get("label") or "").lstrip(":"))
        ]

    indexes = leaked("SHOW INDEX INFO")
    constraints = leaked("SHOW CONSTRAINT INFO")

    def name(row: dict[str, Any]) -> str:
        return (row.get("label") or "").lstrip(":")

    # Nodes are counted from the index rows only — a constraint row carries no
    # count, and a vector-indexed label may have neither.
    empty: dict[str, int] = {}
    counts = reduce(
        lambda acc, row: acc
        | {name(row): acc.get(name(row), 0) + int(row.get("count") or 0)},
        indexes,
        empty,
    )
    zeros = dict.fromkeys(
        [*map(name, constraints), *vector_indexes_by_label(session)], 0
    )
    return {**zeros, **counts}, len(indexes) + len(constraints)


def vector_indexes_by_label(session: Any) -> dict[str, str]:
    """label -> vector index name, for the leaked labels only.

    Vector indexes are a SEPARATE namespace: they do not answer to
    ``DROP INDEX ON :label(prop)`` and their name is not derivable from the
    label (fixtures choose it independently — ``ArchitectureChunk__TEST_x``
    is served by ``architecture_chunk_idx__TEST_x``). The first pass of this
    purge dropped 823 labels and left exactly one entry behind for this
    reason.
    """
    out: dict[str, str] = {}
    for row in session.run("CALL vector_search.show_index_info() YIELD * RETURN *"):
        label = (row["label"] or "").lstrip(":")
        if is_leaked_label(label):
            out[label] = row["index_name"]
    return out


def _as_props(raw: Any) -> list[str]:
    """A row's property field as a list, whatever shape the server used."""
    if isinstance(raw, str):
        return [raw]
    if isinstance(raw, (list, tuple)):
        return [str(p) for p in raw]
    return []


def constrained_properties(session: Any, label: str) -> list[list[str]]:
    """Property TUPLES a uniqueness constraint guards on ``label``.

    Each entry is one constraint, so a composite constraint stays one unit:
    ``(a, b)`` must be dropped as ``ASSERT c.a, c.b IS UNIQUE``. Splitting it
    into two single-property drops removes nothing and raises nothing —
    measured 2026-08-06, when three ``CanonTest_test_*`` constraints on
    ``[a, b]`` survived a purge that reported success.

    Assuming ``id`` is the same class of error: the leaked labels are also
    constrained on ``value`` and ``mkey``.
    """
    rows = map(dict, session.run("SHOW CONSTRAINT INFO"))
    mine = filter(lambda r: (r.get("label") or "").lstrip(":") == label, rows)
    groups = [
        props for props in map(lambda r: _as_props(r.get("properties")), mine) if props
    ]
    return groups or [["id"]]


def indexed_properties(session: Any, label: str) -> list[str]:
    """Properties carrying a label+property index on ``label``.

    ``DROP INDEX ON :label(id)`` is a guess; the leaked labels carry
    ``(value)`` indexes that the guess left behind.
    """
    rows = map(dict, session.run("SHOW INDEX INFO"))
    mine = filter(lambda r: (r.get("label") or "").lstrip(":") == label, rows)
    return [prop for row in mine for prop in _as_props(row.get("property"))]


def drop_label(session: Any, label: str, vector_index: str | None = None) -> None:
    """Delete the label's residual nodes, then all of its schema.

    Nodes first: Memgraph refuses to drop a uniqueness constraint that is
    still backing live nodes, and a leaked label's nodes are test residue by
    construction (the label embeds a uuid from a run that has long finished).
    """
    session.run(f"MATCH (n:`{label}`) DETACH DELETE n").consume()
    statements = [
        "DROP CONSTRAINT ON (c:`{label}`) ASSERT {props} IS UNIQUE".format(
            label=label, props=", ".join(f"c.`{prop}`" for prop in group)
        )
        for group in constrained_properties(session, label)
    ]
    statements += [
        f"DROP INDEX ON :`{label}`({prop})"
        for prop in dict.fromkeys([*indexed_properties(session, label), "id"])
    ]
    statements.append(f"DROP INDEX ON :`{label}`")
    if vector_index:
        statements.append(f"DROP VECTOR INDEX {vector_index}")
    for statement in statements:
        try:
            session.run(statement).consume()
        except Exception as exc:  # noqa: BLE001 - absent DDL is the goal state
            if "doesn't exist" not in str(exc) and "not exist" not in str(exc):
                raise


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--apply", action="store_true", help="drop (default: dry-run report)"
    )
    parser.add_argument(
        "--uri",
        default=None,
        help=f"target bolt/neo4j URI (default: first set of {', '.join(_URI_KEYS)})",
    )
    parser.add_argument(
        "--all-databases",
        action="store_true",
        help=(
            "sweep every database on the instance, not just the default one "
            "(the leak spreads across per-worker databases)"
        ),
    )
    return parser


def _databases(driver: Any) -> list[str]:
    """Every database on the instance, or the default alone if unsupported."""
    with driver.session() as session:
        return sorted(row[0] for row in session.run("SHOW DATABASES"))


def _purge_one(session: Any, label_note: str, apply: bool) -> tuple[int, int]:
    """Survey (and optionally purge) one database. Returns (labels, remaining)."""
    labels, ddl = survey(session)
    residual_nodes = sum(labels.values())
    if not labels:
        return 0, 0
    print(
        f"{label_note}leaked per-test labels: {len(labels)} "
        f"({ddl} index/constraint entries, {residual_nodes} residual nodes)"
    )
    if not apply:
        for label in sorted(labels)[:10]:
            print(f"{label_note}  WOULD drop :{label} ({labels[label]} nodes)")
        if len(labels) > 10:
            print(f"{label_note}  ... and {len(labels) - 10} more")
        return len(labels), 0

    vector_indexes = vector_indexes_by_label(session)
    for label in sorted(labels):
        drop_label(session, label, vector_indexes.get(label))
    remaining, _ = survey(session)
    print(f"{label_note}dropped {len(labels)} labels; remaining: {len(remaining)}")
    return len(labels), len(remaining)


def main(argv: list[str] | None = None, env: Mapping[str, str] | None = None) -> int:
    args = _build_parser().parse_args(argv)
    environ = env if env is not None else os.environ
    uri = resolve_uri(environ, args.uri)

    from neo4j import GraphDatabase

    driver = GraphDatabase.driver(uri, auth=None)
    try:
        print(f"target: {uri}")
        if not args.all_databases:
            with driver.session() as session:
                found, remaining = _purge_one(session, "", args.apply)
            if not found:
                print("clean — nothing to purge")
                return 0
            if not args.apply:
                print("dry run — pass --apply to drop")
            return 0 if not remaining else 1

        total_found = 0
        total_remaining = 0
        for database in _databases(driver):
            with driver.session(database=database) as session:
                found, remaining = _purge_one(session, f"[{database}] ", args.apply)
            total_found += found
            total_remaining += remaining
        if not total_found:
            print("clean — nothing to purge in any database")
            return 0
        print(f"total leaked labels: {total_found}; remaining: {total_remaining}")
        if not args.apply:
            print("dry run — pass --apply to drop")
        return 0 if not total_remaining else 1
    finally:
        driver.close()


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