perf: Improving performance of adding memberships from records

.github/workflows/CI.yaml

@@ -73,6 +73,56 @@ jobs:
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
 
+  benchmarks:
+    name: Benchmarks
+    needs: pre-commit
+    runs-on: ubuntu-latest
+    continue-on-error: true
+    permissions:
+      contents: read
+      pull-requests: write
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          version: ${{ env.UV_VERSION }}
+
+      - name: Set up Python 3.13
+        run: uv python install 3.13
+
+      - name: Install dependencies
+        run: uv sync --all-groups
+
+      - name: Run benchmarks
+        run: uv run pytest benchmarks/ -k "not xlarge_300k" --benchmark-only --benchmark-json=benchmark-results.json --no-cov
+
+      - name: Download previous benchmark data
+        uses: actions/cache@v5
+        with:
+          path: ./cache
+          key: ${{ runner.os }}-benchmark-${{ github.head_ref || github.ref_name }}-${{ github.run_id }}
+          restore-keys: |
+            ${{ runner.os }}-benchmark-${{ github.head_ref || github.ref_name }}-
+            ${{ runner.os }}-benchmark-main-
+
+      - name: Initialize benchmark data file
+        run: mkdir -p cache && test -f cache/benchmark-data.json || echo '[]' > cache/benchmark-data.json
+
+      - name: Publish benchmark results
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          tool: pytest
+          output-file-path: benchmark-results.json
+          external-data-json-path: ./cache/benchmark-data.json
+          fail-on-alert: false
+          comment-on-alert: true
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          alert-threshold: "200%"
+          summary-always: true
+
   package:
     name: Package smoke test
     needs: pytest

.gitignore

@@ -135,6 +135,8 @@ htmlcov/
 .cache
 nosetests.xml
 coverage.xml
+coverage.json
+benchmark-results.json
 *.cover
 *.py,cover
 .hypothesis/

benchmarks/README.md

@@ -0,0 +1,21 @@
+# Benchmark Tests
+
+This folder contains `pytest-benchmark` performance tests for hot paths.
+
+Run all benchmark tests:
+
+```bash
+uv run pytest benchmarks --benchmark-only --no-cov
+```
+
+Run only membership bulk insert benchmark:
+
+```bash
+uv run pytest benchmarks/test_add_memberships_from_records.py --benchmark-only --no-cov
+```
+
+Run the `300k` case only:
+
+```bash
+uv run pytest benchmarks/test_add_memberships_from_records.py -k xlarge_300k --benchmark-only --no-cov
+```

benchmarks/conftest.py

@@ -0,0 +1,99 @@
+"""Shared fixtures for benchmark tests."""
+
+from __future__ import annotations
+
+import uuid
+from collections.abc import Generator
+
+import pytest
+
+from plexosdb import PlexosDB
+
+
+@pytest.fixture(scope="function")
+def db_instance_with_schema() -> Generator[PlexosDB, None, None]:
+    """Create a minimal schema-backed database for benchmark runs."""
+    db = PlexosDB()
+    db.create_schema()
+    with db._db.transaction():
+        db._db.execute(
+            "INSERT INTO t_class(class_id, name, description) VALUES (1, 'System', 'System class')"
+        )
+        db._db.execute(
+            "INSERT INTO t_class(class_id, name, description) VALUES (2, 'Generator', 'Generator class')"
+        )
+        db._db.execute("INSERT INTO t_class(class_id, name, description) VALUES (3, 'Node', 'Node class')")
+        db._db.execute(
+            "INSERT INTO t_class(class_id, name, description) VALUES (4, 'Scenario', 'Scenario class')"
+        )
+        db._db.execute(
+            "INSERT INTO t_class(class_id, name, description) VALUES (5, 'DataFile', 'DataFile class')"
+        )
+        db._db.execute(
+            "INSERT INTO t_class(class_id, name, description) VALUES (6, 'Storage', 'Storage class')"
+        )
+        db._db.execute(
+            "INSERT INTO t_class(class_id, name, description) VALUES (7, 'Report', 'Report class')"
+        )
+        db._db.execute("INSERT INTO t_class(class_id, name, description) VALUES (8, 'Model', 'Model class')")
+        db._db.execute(
+            "INSERT INTO t_object(object_id, name, class_id, GUID) VALUES (1, 'System', 1, ?)",
+            (str(uuid.uuid4()),),
+        )
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (1, 1, 2, 'Generators')"
+        )
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (2, 1, 3, 'Nodes')"
+        )
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (3, 2, 3, 'Nodes')"
+        )
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (4, 1, 4, 'Scenarios')"
+        )
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (5, 1, 6, 'Storages')"
+        )
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (6, 1, 8, 'Models')"
+        )
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (7, 8, 7, 'Models')"
+        )
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (8, 1, 7, 'Reports')"
+        )
+        db._db.execute("INSERT INTO t_unit(unit_id, value) VALUES (1,'MW')")
+        db._db.execute("INSERT INTO t_unit(unit_id, value) VALUES (2,'MWh')")
+        db._db.execute("INSERT INTO t_unit(unit_id, value) VALUES (3,'%')")
+        db._db.execute(
+            "INSERT INTO t_collection(collection_id, parent_class_id, child_class_id, name) "
+            "VALUES (?, ?, ?, ?)",
+            (9, 8, 4, "Scenarios"),
+        )
+        db._db.execute(
+            "INSERT INTO t_property(property_id, collection_id, unit_id, name) VALUES (1,1,1, 'Max Capacity')"
+        )
+        db._db.execute(
+            "INSERT INTO t_property(property_id, collection_id, unit_id, name) VALUES (2,1,2, 'Max Energy')"
+        )
+        db._db.execute(
+            "INSERT INTO t_property(property_id, collection_id, unit_id, name) "
+            "VALUES (3,1,1, 'Rating Factor')"
+        )
+        db._db.execute("INSERT INTO t_config(element, value) VALUES ('Version', '10.0')")
+        db._db.execute("INSERT INTO t_attribute(attribute_id, class_id, name) VALUES( 1, 2, 'Latitude')")
+        db._db.execute(
+            "INSERT INTO t_property_report(property_id, collection_id, name) VALUES (1, 1, 'Units')"
+        )
+    yield db
+    db._db.close()

benchmarks/test_add_memberships_from_records.py

@@ -0,0 +1,97 @@
+"""Benchmark coverage for bulk membership insertion."""
+
+from __future__ import annotations
+
+import uuid
+
+import pytest
+
+from plexosdb import ClassEnum, CollectionEnum, PlexosDB
+
+
+def _insert_objects(
+    db: PlexosDB,
+    *,
+    class_id: int,
+    count: int,
+    prefix: str,
+    start_id: int,
+) -> list[int]:
+    object_ids = [start_id + idx for idx in range(count)]
+    params = [
+        (object_id, f"{prefix}_{idx}", class_id, str(uuid.uuid4()))
+        for idx, object_id in enumerate(object_ids)
+    ]
+    db._db.executemany("INSERT INTO t_object(object_id, name, class_id, GUID) VALUES (?, ?, ?, ?)", params)
+    return object_ids
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    ("record_count", "chunksize", "rounds"),
+    [
+        pytest.param(100, 100, 10, id="small"),
+        pytest.param(1_000, 1_000, 10, id="medium"),
+        pytest.param(10_000, 10_000, 10, id="large"),
+        pytest.param(300_000, 10_000, 2, id="xlarge_300k"),
+    ],
+)
+def test_add_memberships_from_records_benchmark(
+    benchmark,
+    db_instance_with_schema: PlexosDB,
+    record_count: int,
+    chunksize: int,
+    rounds: int,
+) -> None:
+    """Benchmark `add_memberships_from_records` across different payload sizes."""
+    db = db_instance_with_schema
+    parent_class_id = db.get_class_id(ClassEnum.Generator)
+    child_class_id = db.get_class_id(ClassEnum.Node)
+    collection_id = db.get_collection_id(
+        CollectionEnum.Nodes,
+        parent_class_enum=ClassEnum.Generator,
+        child_class_enum=ClassEnum.Node,
+    )
+    parent_ids = _insert_objects(
+        db,
+        class_id=parent_class_id,
+        count=1,
+        prefix=f"benchmark_parent_{record_count}",
+        start_id=10_000,
+    )
+    child_ids = _insert_objects(
+        db,
+        class_id=child_class_id,
+        count=record_count,
+        prefix=f"benchmark_child_{record_count}",
+        start_id=50_000,
+    )
+    records = [
+        {
+            "parent_class_id": parent_class_id,
+            "parent_object_id": parent_ids[0],
+            "collection_id": collection_id,
+            "child_class_id": child_class_id,
+            "child_object_id": child_id,
+        }
+        for child_id in child_ids
+    ]
+
+    def _reset_memberships() -> None:
+        db._db.execute(
+            (
+                "DELETE FROM t_membership "
+                "WHERE collection_id = ? AND parent_class_id = ? AND child_class_id = ?"
+            ),
+            (collection_id, parent_class_id, child_class_id),
+        )
+
+    result = benchmark.pedantic(
+        db.add_memberships_from_records,
+        args=(records,),
+        kwargs={"chunksize": chunksize},
+        setup=_reset_memberships,
+        rounds=rounds,
+        iterations=1,
+    )
+    assert result is True