cleveragents-core/benchmarks/project_repository_bench.py

"""ASV benchmarks for project repository operations.

Measures:
- NamespacedProjectRepository create/get/list throughput
- ProjectResourceLinkRepository create_link/list_links throughput
"""

from __future__ import annotations

from datetime import UTC, datetime

from sqlalchemy import create_engine, event, text
from sqlalchemy.orm import sessionmaker

from cleveragents.domain.models.core.project import NamespacedProject
from cleveragents.infrastructure.database.models import (
    Base,
    ResourceModel,
    ResourceTypeModel,
)
from cleveragents.infrastructure.database.repositories import (
    NamespacedProjectRepository,
    ProjectResourceLinkRepository,
)


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


def _set_sqlite_pragma(dbapi_conn: object, _connection_record: object) -> None:
    cursor = dbapi_conn.cursor()  # type: ignore[union-attr]
    cursor.execute("PRAGMA foreign_keys = ON")
    cursor.close()


class ProjectRepositoryCRUD:
    """Benchmark project repository CRUD operations."""

    def setup(self) -> None:
        self.engine = create_engine("sqlite:///:memory:")
        event.listen(self.engine, "connect", _set_sqlite_pragma)
        with self.engine.connect() as conn:
            conn.execute(text("PRAGMA foreign_keys = ON"))
            conn.commit()
        Base.metadata.create_all(self.engine)
        self.sf = sessionmaker(bind=self.engine)

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

    def time_create_and_get_project(self) -> None:
        repo = NamespacedProjectRepository(session_factory=self.sf)
        session = self.sf()
        project = NamespacedProject(name="bench-proj", namespace="bench")
        repo.create(project)
        session.commit()
        repo.get("bench/bench-proj")
        repo.delete("bench/bench-proj")
        session.commit()

    def time_list_50_projects(self) -> None:
        repo = NamespacedProjectRepository(session_factory=self.sf)
        session = self.sf()
        for i in range(50):
            repo.create(NamespacedProject(name=f"proj-{i}", namespace="bench"))
            session.commit()
        repo.list_projects(namespace="bench", limit=50)
        for i in range(50):
            repo.delete(f"bench/proj-{i}")
            session.commit()


class LinkRepositoryCRUD:
    """Benchmark link repository operations."""

    def setup(self) -> None:
        self.engine = create_engine("sqlite:///:memory:")
        event.listen(self.engine, "connect", _set_sqlite_pragma)
        with self.engine.connect() as conn:
            conn.execute(text("PRAGMA foreign_keys = ON"))
            conn.commit()
        Base.metadata.create_all(self.engine)

        # Use a single session for all setup operations so FK
        # relationships are visible within the same transaction.
        self._setup_session = sessionmaker(bind=self.engine)()
        self.sf = sessionmaker(bind=self.engine)

        session = self._setup_session
        # Resource type
        rt = ResourceTypeModel()
        rt.name = "bench/git-checkout"  # type: ignore[assignment]
        rt.namespace = "bench"  # type: ignore[assignment]
        rt.resource_kind = "physical"  # type: ignore[assignment]
        rt.user_addable = True  # type: ignore[assignment]
        rt.created_at = _now_iso()  # type: ignore[assignment]
        rt.updated_at = _now_iso()  # type: ignore[assignment]
        session.add(rt)
        session.commit()

        # Project — use the same session so the commit is on the same
        # connection that created the resource type.
        project_repo = NamespacedProjectRepository(
            session_factory=lambda: session,
        )
        project_repo.create(NamespacedProject(name="link-bench", namespace="bench"))
        session.commit()

        # Resources
        for i in range(30):
            r = ResourceModel()
            r.resource_id = str(i).zfill(26)  # type: ignore[assignment]
            r.type_name = "bench/git-checkout"  # type: ignore[assignment]
            r.resource_kind = "physical"  # type: ignore[assignment]
            r.created_at = _now_iso()  # type: ignore[assignment]
            r.updated_at = _now_iso()  # type: ignore[assignment]
            session.add(r)
        session.commit()

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

    def time_create_and_list_30_links(self) -> None:
        # Use a scoped session so create_link and remove_link share state.
        session = self.sf()
        link_repo = ProjectResourceLinkRepository(session_factory=lambda: session)
        link_ids: list[str] = []
        for i in range(30):
            link = link_repo.create_link(
                project_name="bench/link-bench",
                resource_id=str(i).zfill(26),
            )
            session.commit()
            link_ids.append(str(link.link_id))
        link_repo.list_links("bench/link-bench")
        for lid in link_ids:
            link_repo.remove_link(lid)
            session.commit()
        session.close()