cleveragents-core/benchmarks/db_resource_bench.py

"""ASV benchmarks for database resource types and transaction sandbox.

Measures the performance of:
- Database type definition lookup
- Connection argument resolution
- Connection validation (SQLite)
- TransactionSandbox create/commit/rollback lifecycle
- SandboxFactory routing for transaction_rollback
"""

from __future__ import annotations

import importlib
import os
import sqlite3
import sys
import tempfile
from pathlib import Path

# Ensure the local *source* tree is importable even when ASV has an
# older build of the package installed.
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

# Force-reload so ASV picks up the source tree version.
import cleveragents  # noqa: E402

importlib.reload(cleveragents)

from cleveragents.infrastructure.sandbox.factory import (  # noqa: E402
    SandboxFactory,
)
from cleveragents.infrastructure.sandbox.transaction_sandbox import (  # noqa: E402
    TransactionSandbox,
)
from cleveragents.resource.handlers.database import (  # noqa: E402
    DATABASE_TYPE_DEFS,
    DatabaseResourceHandler,
    resolve_connection_args,
    validate_connection,
)

# ---------------------------------------------------------------------------
# Type definition benchmarks
# ---------------------------------------------------------------------------


class TypeDefinitionSuite:
    """Benchmark database type definition operations."""

    timeout = 60

    def time_lookup_all_type_defs(self) -> None:
        """Time iterating all database type definitions."""
        for td in DATABASE_TYPE_DEFS:
            _ = td["name"]

    def time_resolve_postgres_args(self) -> None:
        """Time resolving postgres connection args."""
        resolve_connection_args(
            {"host": "db.example.com", "port": 5432, "dbname": "mydb"},
            "postgres",
        )

    def time_resolve_sqlite_args(self) -> None:
        """Time resolving sqlite connection args."""
        resolve_connection_args({"path": "/tmp/test.db"}, "sqlite")

    def time_handler_instantiation(self) -> None:
        """Time creating a DatabaseResourceHandler."""
        DatabaseResourceHandler()


# ---------------------------------------------------------------------------
# Connection validation benchmarks
# ---------------------------------------------------------------------------


class ConnectionValidationSuite:
    """Benchmark connection validation operations."""

    timeout = 60

    def setup(self) -> None:
        """Create a temp SQLite database for benchmarking."""
        fd, self.db_path = tempfile.mkstemp(suffix=".db")
        os.close(fd)
        conn = sqlite3.connect(self.db_path)
        conn.execute("CREATE TABLE bench (id INTEGER PRIMARY KEY)")
        conn.commit()
        conn.close()

    def teardown(self) -> None:
        """Remove the temp database."""
        if hasattr(self, "db_path") and os.path.exists(self.db_path):
            os.unlink(self.db_path)

    def time_validate_sqlite_file(self) -> None:
        """Time validating a file-based SQLite connection."""
        validate_connection("sqlite", {"path": self.db_path})

    def time_validate_sqlite_memory(self) -> None:
        """Time validating an in-memory SQLite connection."""
        validate_connection("sqlite", {"path": ":memory:"})

    def time_validate_postgres_stub(self) -> None:
        """Time validating a postgres stub (no driver)."""
        validate_connection(
            "postgres",
            {"host": "localhost", "port": 5432, "dbname": "bench"},
        )


# ---------------------------------------------------------------------------
# Transaction sandbox benchmarks
# ---------------------------------------------------------------------------


class TransactionSandboxSuite:
    """Benchmark TransactionSandbox lifecycle operations."""

    timeout = 60

    def setup(self) -> None:
        """Create a temp SQLite database for benchmarking."""
        fd, self.db_path = tempfile.mkstemp(suffix=".db")
        os.close(fd)
        conn = sqlite3.connect(self.db_path)
        conn.execute("CREATE TABLE bench (id INTEGER PRIMARY KEY, val TEXT)")
        conn.commit()
        conn.close()

    def teardown(self) -> None:
        """Remove the temp database."""
        if hasattr(self, "db_path") and os.path.exists(self.db_path):
            os.unlink(self.db_path)

    def time_create_and_cleanup(self) -> None:
        """Time sandbox create + cleanup cycle."""
        sandbox = TransactionSandbox(resource_id="bench-1", original_path=self.db_path)
        sandbox.create(plan_id="BENCH-PLAN")
        sandbox.cleanup()

    def time_create_insert_commit(self) -> None:
        """Time full sandbox lifecycle: create, insert, commit, cleanup."""
        sandbox = TransactionSandbox(resource_id="bench-2", original_path=self.db_path)
        sandbox.create(plan_id="BENCH-PLAN")
        sandbox.execute("INSERT INTO bench (id, val) VALUES (1, 'bench')")
        sandbox.commit("bench commit")
        sandbox.cleanup()

    def time_create_insert_rollback(self) -> None:
        """Time sandbox lifecycle with rollback."""
        sandbox = TransactionSandbox(resource_id="bench-3", original_path=self.db_path)
        sandbox.create(plan_id="BENCH-PLAN")
        sandbox.execute("INSERT INTO bench (id, val) VALUES (2, 'rollback')")
        sandbox.rollback()
        sandbox.cleanup()


# ---------------------------------------------------------------------------
# Factory routing benchmarks
# ---------------------------------------------------------------------------


class FactoryRoutingSuite:
    """Benchmark SandboxFactory routing for database strategies."""

    timeout = 60

    def time_factory_transaction_rollback(self) -> None:
        """Time factory creating a TransactionSandbox."""
        factory = SandboxFactory()
        factory.create_sandbox(
            resource_id="bench-factory",
            original_path=":memory:",
            sandbox_strategy="transaction_rollback",
        )