cleveragents-core/benchmarks/resource_repository_bench.py

"""ASV benchmarks for resource repository operations."""

from datetime import UTC, datetime


def _setup_db():
    """Create an in-memory database with schema and return session factory."""
    from sqlalchemy import create_engine, event
    from sqlalchemy.orm import sessionmaker

    from cleveragents.infrastructure.database.models import Base

    engine = create_engine("sqlite:///:memory:")

    @event.listens_for(engine, "connect")
    def _fk(conn, _rec):
        conn.cursor().execute("PRAGMA foreign_keys=ON")

    Base.metadata.create_all(engine)
    return sessionmaker(bind=engine)


def _seed_type(rt_repo, name="bench/seed-type"):
    """Seed a resource type for benchmarks."""
    from cleveragents.domain.models.core.resource_type import (
        ResourceKind,
        ResourceTypeSpec,
        SandboxStrategy,
    )

    spec = ResourceTypeSpec(
        name=name,
        description="Bench type",
        resource_kind=ResourceKind.PHYSICAL,
        sandbox_strategy=SandboxStrategy.NONE,
        user_addable=True,
        cli_args=[],
        parent_types=[],
        child_types=[],
        auto_discovery=None,
        equivalence=None,
        handler=None,
        capabilities={
            "read": True,
            "write": True,
            "sandbox": True,
            "checkpoint": False,
        },
        built_in=False,
    )
    rt_repo.create(spec)
    return spec


_CB32 = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
_BENCH_CTR = 0


def _bench_ulid():
    global _BENCH_CTR
    _BENCH_CTR += 1
    n = _BENCH_CTR
    suffix = ""
    for _ in range(8):
        suffix = _CB32[n % 32] + suffix
        n //= 32
    return f"01HGZ6FE0AQDYTR4BX{suffix}"


class TimeResourceTypeCreate:
    """Benchmark ResourceTypeRepository.create() throughput."""

    def setup(self):
        from cleveragents.infrastructure.database.repositories import (
            ResourceTypeRepository,
        )

        self.factory = _setup_db()
        self.repo = ResourceTypeRepository(self.factory)
        self._counter = 0

    def time_create_resource_type(self):
        from cleveragents.domain.models.core.resource_type import (
            ResourceKind,
            ResourceTypeSpec,
            SandboxStrategy,
        )

        self._counter += 1
        spec = ResourceTypeSpec(
            name=f"bench/type-{self._counter}",
            description="Bench",
            resource_kind=ResourceKind.PHYSICAL,
            sandbox_strategy=SandboxStrategy.NONE,
            user_addable=True,
            cli_args=[],
            parent_types=[],
            child_types=[],
            auto_discovery=None,
            equivalence=None,
            handler=None,
            capabilities={
                "read": True,
                "write": True,
                "sandbox": True,
                "checkpoint": False,
            },
            built_in=False,
        )
        self.repo.create(spec)


class TimeResourceTypeList:
    """Benchmark ResourceTypeRepository.list_types() throughput."""

    def setup(self):
        from cleveragents.infrastructure.database.repositories import (
            ResourceTypeRepository,
        )

        self.factory = _setup_db()
        self.repo = ResourceTypeRepository(self.factory)
        for i in range(50):
            _seed_type(self.repo, name=f"bench/list-type-{i}")

    def time_list_all_types(self):
        self.repo.list_types()

    def time_list_by_namespace(self):
        self.repo.list_types(namespace="bench")


class TimeResourceCreate:
    """Benchmark ResourceRepository.create() throughput."""

    def setup(self):
        from cleveragents.infrastructure.database.repositories import (
            ResourceRepository,
            ResourceTypeRepository,
        )

        self.factory = _setup_db()
        self._session = self.factory()
        shared = lambda: self._session  # noqa: E731
        self.rt_repo = ResourceTypeRepository(shared)
        self.res_repo = ResourceRepository(shared)
        _seed_type(self.rt_repo)
        self._session.commit()
        self._counter = 0

    def time_create_resource(self):
        from cleveragents.domain.models.core.resource import (
            PhysVirt,
            Resource,
            ResourceCapabilities,
        )

        self._counter += 1
        res = Resource(
            resource_id=_bench_ulid(),
            name=f"bench/res-{self._counter}",
            resource_type_name="bench/seed-type",
            classification=PhysVirt.PHYSICAL,
            description="Bench resource",
            properties={},
            location="/tmp/bench",
            content_hash=None,
            sandbox_strategy=None,
            capabilities=ResourceCapabilities(),
            created_at=datetime.now(tz=UTC),
            updated_at=datetime.now(tz=UTC),
        )
        self.res_repo.create(res)


class TimeResourceResolve:
    """Benchmark ResourceRepository.resolve_namespaced_name() throughput."""

    def setup(self):
        from cleveragents.domain.models.core.resource import (
            PhysVirt,
            Resource,
            ResourceCapabilities,
        )
        from cleveragents.infrastructure.database.repositories import (
            ResourceRepository,
            ResourceTypeRepository,
        )

        self.factory = _setup_db()
        self._session = self.factory()
        shared = lambda: self._session  # noqa: E731
        self.rt_repo = ResourceTypeRepository(shared)
        self.res_repo = ResourceRepository(shared)
        _seed_type(self.rt_repo)
        self._session.commit()
        self.ulid = _bench_ulid()
        res = Resource(
            resource_id=self.ulid,
            name="bench/resolve-target",
            resource_type_name="bench/seed-type",
            classification=PhysVirt.PHYSICAL,
            description="Resolve target",
            properties={},
            location="/tmp/bench",
            content_hash=None,
            sandbox_strategy=None,
            capabilities=ResourceCapabilities(),
            created_at=datetime.now(tz=UTC),
            updated_at=datetime.now(tz=UTC),
        )
        self.res_repo.create(res)
        self._session.commit()

    def time_resolve_by_name(self):
        self.res_repo.resolve_namespaced_name("bench/resolve-target")

    def time_resolve_by_ulid(self):
        self.res_repo.resolve_namespaced_name(self.ulid)