cleveragents-core/benchmarks/resource_registry_migration_bench.py

"""ASV benchmarks for resource registry migration and ORM operations.

Measures:
- Schema creation time (Base.metadata.create_all)
- ResourceTypeModel insert throughput
- ResourceModel insert throughput
- ResourceEdgeModel insert throughput
- DAG query performance (children/parents lookup)
"""

from __future__ import annotations

import json
from datetime import UTC, datetime

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from cleveragents.infrastructure.database.models import (
    Base,
    ResourceEdgeModel,
    ResourceModel,
    ResourceTypeModel,
)


def _now_iso() -> str:
    return datetime.now(tz=UTC).isoformat()


def _make_ulid(n: int) -> str:
    return str(n).zfill(26)


class SchemaCreation:
    """Benchmark schema creation (all tables via metadata.create_all)."""

    def time_create_all_tables(self) -> None:
        engine = create_engine("sqlite:///:memory:")
        Base.metadata.create_all(engine)
        engine.dispose()


class ResourceTypeInsert:
    """Benchmark inserting resource type records."""

    def setup(self) -> None:
        self.engine = create_engine("sqlite:///:memory:")
        Base.metadata.create_all(self.engine)
        self.session_factory = sessionmaker(bind=self.engine)
        self._iter = 0

    def teardown(self) -> None:
        self.engine.dispose()

    def time_insert_single_resource_type(self) -> None:
        self._iter += 1
        session = self.session_factory()
        rt = ResourceTypeModel()
        rt.name = f"bench/type-single-{self._iter}"
        rt.namespace = "bench"
        rt.resource_kind = "physical"
        rt.user_addable = True
        rt.created_at = _now_iso()
        rt.updated_at = _now_iso()
        session.add(rt)
        session.commit()
        session.close()

    def time_insert_100_resource_types(self) -> None:
        self._iter += 1
        batch = self._iter
        session = self.session_factory()
        for i in range(100):
            rt = ResourceTypeModel()
            rt.name = f"bench/type-{batch}-{i}"
            rt.namespace = "bench"
            rt.resource_kind = "physical" if i % 2 == 0 else "virtual"
            rt.user_addable = i % 3 == 0
            rt.args_schema_json = json.dumps([{"flag": f"--arg-{batch}-{i}"}])
            rt.created_at = _now_iso()
            rt.updated_at = _now_iso()
            session.add(rt)
        session.commit()
        session.close()


class ResourceInsert:
    """Benchmark inserting resource records."""

    def setup(self) -> None:
        self.engine = create_engine("sqlite:///:memory:")
        Base.metadata.create_all(self.engine)
        self.session_factory = sessionmaker(bind=self.engine)
        session = self.session_factory()
        rt = ResourceTypeModel()
        rt.name = "bench/git-checkout"
        rt.namespace = "bench"
        rt.resource_kind = "physical"
        rt.user_addable = True
        rt.created_at = _now_iso()
        rt.updated_at = _now_iso()
        session.add(rt)
        session.commit()
        session.close()
        self._iter = 0

    def teardown(self) -> None:
        self.engine.dispose()

    def time_insert_100_resources(self) -> None:
        self._iter += 1
        batch = self._iter
        session = self.session_factory()
        for i in range(100):
            idx = batch * 1000 + i
            r = ResourceModel()
            r.resource_id = _make_ulid(idx)
            r.namespaced_name = f"bench/resource-{batch}-{i}"
            r.namespace = "bench"
            r.type_name = "bench/git-checkout"
            r.resource_kind = "physical"
            r.location = f"/path/to/resource-{idx}"
            r.properties_json = json.dumps({"index": idx})
            r.created_at = _now_iso()
            r.updated_at = _now_iso()
            session.add(r)
        session.commit()
        session.close()


class ResourceEdgeQuery:
    """Benchmark DAG edge queries (children/parents lookup)."""

    def setup(self) -> None:
        self.engine = create_engine("sqlite:///:memory:")
        Base.metadata.create_all(self.engine)
        self.session_factory = sessionmaker(bind=self.engine)
        session = self.session_factory()

        # Create types
        for tname in ("bench/parent-type", "bench/child-type"):
            rt = ResourceTypeModel()
            rt.name = tname
            rt.namespace = "bench"
            rt.resource_kind = "physical"
            rt.user_addable = True
            rt.created_at = _now_iso()
            rt.updated_at = _now_iso()
            session.add(rt)

        # Create 1 parent + 50 children with edges
        parent = ResourceModel()
        parent.resource_id = _make_ulid(0)
        parent.namespaced_name = "bench/parent"
        parent.namespace = "bench"
        parent.type_name = "bench/parent-type"
        parent.resource_kind = "physical"
        parent.created_at = _now_iso()
        parent.updated_at = _now_iso()
        session.add(parent)
        self.parent_id = parent.resource_id

        for i in range(1, 51):
            child = ResourceModel()
            child.resource_id = _make_ulid(i)
            child.type_name = "bench/child-type"
            child.resource_kind = "physical"
            child.created_at = _now_iso()
            child.updated_at = _now_iso()
            session.add(child)

            edge = ResourceEdgeModel()
            edge.parent_id = parent.resource_id
            edge.child_id = child.resource_id
            edge.link_type = "contains"
            edge.created_at = _now_iso()
            session.add(edge)

        session.commit()
        session.close()

    def teardown(self) -> None:
        self.engine.dispose()

    def time_query_50_children(self) -> None:
        session = self.session_factory()
        edges = (
            session.query(ResourceEdgeModel)
            .filter(ResourceEdgeModel.parent_id == self.parent_id)
            .all()
        )
        assert len(edges) == 50
        session.close()

    def time_query_parent_of_child(self) -> None:
        session = self.session_factory()
        child_id = _make_ulid(25)
        edges = (
            session.query(ResourceEdgeModel)
            .filter(ResourceEdgeModel.child_id == child_id)
            .all()
        )
        assert len(edges) == 1
        session.close()