cleveragents-core/benchmarks/project_migration_bench.py

"""ASV benchmarks for project migration and ORM operations.

Measures:
- NamespacedProjectModel insert throughput
- ProjectResourceLinkModel insert throughput
- Project-to-domain conversion performance
- Link query performance (by project name)
"""

from __future__ import annotations

import json
import threading
from datetime import UTC, datetime

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

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

# Thread-safe counter so repeated ASV iterations never collide on names/IDs.
_iter_lock = threading.Lock()
_iter_counter = 0


def _next_iter() -> int:
    """Return a monotonically increasing iteration number."""
    global _iter_counter
    with _iter_lock:
        _iter_counter += 1
        return _iter_counter


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


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


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 ProjectInsert:
    """Benchmark inserting project records."""

    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.session_factory = sessionmaker(bind=self.engine)

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

    def time_insert_single_project(self) -> None:
        it = _next_iter()
        session = self.session_factory()
        proj = NamespacedProjectModel(
            namespaced_name=f"bench/single-project-{it}",
            namespace="bench",
            tags_json="[]",
            created_at=_now_iso(),
            updated_at=_now_iso(),
        )
        session.add(proj)
        session.commit()
        session.close()

    def time_insert_100_projects(self) -> None:
        it = _next_iter()
        session = self.session_factory()
        for i in range(100):
            proj = NamespacedProjectModel(
                namespaced_name=f"bench/project-{it}-{i}",
                namespace="bench",
                description=f"Project {it}-{i}",
                tags_json=json.dumps([f"tag-{it}-{i}"]),
                created_at=_now_iso(),
                updated_at=_now_iso(),
            )
            session.add(proj)
        session.commit()
        session.close()


class ProjectLinkInsert:
    """Benchmark inserting project resource link records."""

    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.session_factory = sessionmaker(bind=self.engine)

        session = self.session_factory()
        # Create resource type
        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)

        # Create resources (shared pool for link targets)
        for i in range(50):
            r = ResourceModel()
            r.resource_id = _make_ulid(i)
            r.type_name = "bench/git-checkout"
            r.resource_kind = "physical"
            r.created_at = _now_iso()
            r.updated_at = _now_iso()
            session.add(r)

        session.commit()
        session.close()

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

    def time_insert_50_links(self) -> None:
        it = _next_iter()
        base = it * 10000
        session = self.session_factory()
        # Create a unique project per iteration so the composite UNIQUE
        # constraint (project_name, resource_id) never collides.
        proj_name = f"bench/link-project-{it}"
        proj = NamespacedProjectModel(
            namespaced_name=proj_name,
            namespace="bench",
            tags_json="[]",
            created_at=_now_iso(),
            updated_at=_now_iso(),
        )
        session.add(proj)
        session.flush()
        for i in range(50):
            link = ProjectResourceLinkModel(
                link_id=_make_ulid(base + i),
                project_name=proj_name,
                resource_id=_make_ulid(i),
                read_only=i % 2 == 0,
                created_at=_now_iso(),
            )
            session.add(link)
        session.commit()
        session.close()


class ProjectDomainConversion:
    """Benchmark domain model conversion."""

    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.session_factory = sessionmaker(bind=self.engine)

        session = self.session_factory()
        proj = NamespacedProjectModel(
            namespaced_name="bench/domain-project",
            namespace="bench",
            description="Domain conversion bench",
            invariants_json=json.dumps(["inv-1", "inv-2", "inv-3"]),
            tags_json=json.dumps(["tag-a", "tag-b"]),
            created_at=_now_iso(),
            updated_at=_now_iso(),
        )
        session.add(proj)
        session.commit()
        session.close()

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

    def time_to_domain_conversion(self) -> None:
        session = self.session_factory()
        proj = session.get(NamespacedProjectModel, "bench/domain-project")
        assert proj is not None
        proj.to_domain()
        session.close()


class ProjectLinkQuery:
    """Benchmark link queries by project name."""

    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.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)

        proj = NamespacedProjectModel(
            namespaced_name="bench/query-project",
            namespace="bench",
            tags_json="[]",
            created_at=_now_iso(),
            updated_at=_now_iso(),
        )
        session.add(proj)

        for i in range(30):
            r = ResourceModel()
            r.resource_id = _make_ulid(i + 2000)
            r.type_name = "bench/git-checkout"
            r.resource_kind = "physical"
            r.created_at = _now_iso()
            r.updated_at = _now_iso()
            session.add(r)

            link = ProjectResourceLinkModel(
                link_id=_make_ulid(i + 3000),
                project_name="bench/query-project",
                resource_id=_make_ulid(i + 2000),
                read_only=False,
                created_at=_now_iso(),
            )
            session.add(link)

        session.commit()
        session.close()

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

    def time_query_30_links_by_project(self) -> None:
        session = self.session_factory()
        links = (
            session.query(ProjectResourceLinkModel)
            .filter(ProjectResourceLinkModel.project_name == "bench/query-project")
            .all()
        )
        assert len(links) == 30
        session.close()