docqube_backend/scripts/bench_subtree.py

"""
C3 — is the materialised path actually faster than the recursive CTE?

A performance change nobody measured is a guess with extra code. This builds a
tree of a given shape, resolves the same subtree both ways, and prints the
difference.

    APP_ENV=rework python scripts/bench_subtree.py --depth 8 --breadth 4

**It writes to the database it is pointed at**, in a transaction it rolls back.
Point it at a scratch database, not production.

The honest expectation: at DocQube's current scale — a handful of departments,
two or three levels — the difference will be small, and this exists to say so
with a number rather than to justify the change after the fact. It matters at
depth, and the number tells you where "depth" begins.
"""

import argparse
import os
import statistics
import sys
import time
import uuid
from pathlib import Path

os.environ.setdefault("APP_ENV", "rework")
os.environ.setdefault("PYTHONIOENCODING", "utf-8")
sys.path.insert(0, str(Path(__file__).parent.parent))


def build_tree(conn, tenant_id, depth: int, breadth: int):
    """A tree of `breadth ** depth` leaves, with paths, returning the root id."""
    from sqlalchemy import text

    root_id = uuid.uuid4()
    conn.execute(
        text(
            "INSERT INTO org_units (id, tenant_id, name, depth, sort_order, "
            "is_deleted, path, created_at, updated_at) "
            "VALUES (:id, :t, 'bench-root', 0, 0, false, :path, now(), now())"
        ),
        {"id": root_id, "t": tenant_id, "path": f"/{root_id}/"},
    )

    level = [(root_id, f"/{root_id}/")]
    total = 1
    for d in range(1, depth + 1):
        nxt = []
        for parent_id, parent_path in level:
            for b in range(breadth):
                node_id = uuid.uuid4()
                node_path = f"{parent_path}{node_id}/"
                conn.execute(
                    text(
                        "INSERT INTO org_units (id, tenant_id, parent_id, name, "
                        "depth, sort_order, is_deleted, path, created_at, updated_at) "
                        "VALUES (:id, :t, :p, :n, :d, 0, false, :path, now(), now())"
                    ),
                    {
                        "id": node_id,
                        "t": tenant_id,
                        "p": parent_id,
                        "n": f"bench-{d}-{b}",
                        "d": d,
                        "path": node_path,
                    },
                )
                nxt.append((node_id, node_path))
                total += 1
        level = nxt
    return root_id, total


def time_it(fn, runs: int) -> float:
    """Median of `runs`, in milliseconds. Median, because one slow run is noise."""
    samples = []
    for _ in range(runs):
        started = time.perf_counter()
        fn()
        samples.append((time.perf_counter() - started) * 1000)
    return statistics.median(samples)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--depth", type=int, default=6)
    parser.add_argument("--breadth", type=int, default=3)
    parser.add_argument("--runs", type=int, default=25)
    args = parser.parse_args()

    from sqlalchemy import create_engine, text

    from app.core.settings import settings

    engine = create_engine(settings.DATABASE_URL)
    conn = engine.connect()
    trans = conn.begin()

    try:
        # Its own tenant, inside the transaction that gets rolled back. Using a
        # real one would make the benchmark depend on seed data and would put
        # tens of thousands of rows under somebody's actual account.
        tenant_id = uuid.uuid4()
        conn.execute(
            text(
                "INSERT INTO tenants (id, name, slug, is_active, is_deleted, "
                "storage_quota_bytes, created_at, updated_at) "
                "VALUES (:id, 'bench', :slug, true, false, 1073741824, now(), now())"
            ),
            {"id": tenant_id, "slug": f"bench-{tenant_id.hex[:8]}"},
        )

        print(f"building depth={args.depth} breadth={args.breadth} ...")
        root_id, total = build_tree(conn, tenant_id, args.depth, args.breadth)
        print(f"{total} units\n")

        root_path = conn.execute(
            text("SELECT path FROM org_units WHERE id = :id"), {"id": root_id}
        ).scalar()

        def by_cte():
            conn.execute(
                text(
                    "WITH RECURSIVE t AS ("
                    "  SELECT id FROM org_units "
                    "   WHERE id = :root AND tenant_id = :tenant AND NOT is_deleted "
                    "  UNION "
                    "  SELECT c.id FROM org_units c JOIN t ON c.parent_id = t.id "
                    "   WHERE c.tenant_id = :tenant AND NOT c.is_deleted"
                    ") SELECT count(*) FROM t"
                ),
                {"root": root_id, "tenant": tenant_id},
            ).scalar()

        def by_path():
            conn.execute(
                text(
                    "SELECT count(*) FROM org_units "
                    " WHERE tenant_id = :tenant AND NOT is_deleted "
                    "   AND path LIKE :prefix"
                ),
                {"tenant": tenant_id, "prefix": f"{root_path}%"},
            ).scalar()

        # Warm both, so the first run's planning does not become the result.
        by_cte()
        by_path()

        cte_ms = time_it(by_cte, args.runs)
        path_ms = time_it(by_path, args.runs)

        print(f"{'recursive CTE':<20}{cte_ms:8.2f} ms")
        print(f"{'path prefix':<20}{path_ms:8.2f} ms")
        if path_ms > 0:
            print(f"{'speedup':<20}{cte_ms / path_ms:8.2f}x")
        return 0
    finally:
        # Always. This inserts a lot of rows and none of them should survive.
        trans.rollback()
        conn.close()
        engine.dispose()


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