benchmarks.py

"""
ORM Benchmarks: TortoiseORM vs SQLAlchemy with asyncpg

This script benchmarks various database operations comparing
TortoiseORM and SQLAlchemy, both using the asyncpg driver.
"""

import asyncio
import statistics
import time
import uuid
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import date
from decimal import Decimal
from typing import Any, Callable, Coroutine

from rich.console import Console
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn
from sqlalchemy import select, update, delete, func as sa_func
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.orm import selectinload
from tortoise import Tortoise
from tortoise.functions import Count, Avg

from config import get_database_url, get_tortoise_config, BENCHMARK_SETTINGS
from models.sqlalchemy_models import Base as SABase
from models.sqlalchemy_models import Author as SAAuthor, Book as SABook, Publisher as SAPublisher, Review as SAReview
from models.tortoise_models import Author as TAuthor, Book as TBook, Publisher as TPublisher, Review as TReview

console = Console()


@dataclass
class BenchmarkResult:
    """Holds the result of a benchmark."""
    name: str
    orm: str
    mean_ms: float
    std_ms: float
    min_ms: float
    max_ms: float
    iterations: int


class BenchmarkRunner:
    """Runs benchmarks and collects results."""

    def __init__(self):
        self.results: list[BenchmarkResult] = []
        self.sa_engine = None
        self.sa_session_factory = None

    async def setup_sqlalchemy(self):
        """Initialize SQLAlchemy engine and session factory."""
        self.sa_engine = create_async_engine(
            get_database_url(),
            echo=False,
            pool_size=10,
            max_overflow=20,
        )
        self.sa_session_factory = async_sessionmaker(
            self.sa_engine,
            class_=AsyncSession,
            expire_on_commit=False
        )

        # Create tables
        async with self.sa_engine.begin() as conn:
            await conn.run_sync(SABase.metadata.create_all)

    async def setup_tortoise(self):
        """Initialize Tortoise ORM."""
        await Tortoise.init(config=get_tortoise_config())
        await Tortoise.generate_schemas()

    async def cleanup_sqlalchemy(self):
        """Cleanup SQLAlchemy resources."""
        if self.sa_engine:
            async with self.sa_engine.begin() as conn:
                await conn.run_sync(SABase.metadata.drop_all)
            await self.sa_engine.dispose()

    async def cleanup_tortoise(self):
        """Cleanup Tortoise ORM resources."""
        # Drop tables
        conn = Tortoise.get_connection("default")
        await conn.execute_script("""
            DROP TABLE IF EXISTS tortoise_reviews CASCADE;
            DROP TABLE IF EXISTS tortoise_books CASCADE;
            DROP TABLE IF EXISTS tortoise_authors CASCADE;
            DROP TABLE IF EXISTS tortoise_publishers CASCADE;
        """)
        await Tortoise.close_connections()

    @asynccontextmanager
    async def sa_session(self):
        """Context manager for SQLAlchemy sessions."""
        async with self.sa_session_factory() as session:
            yield session
            await session.commit()

    async def time_operation(
        self,
        operation: Callable[[], Coroutine[Any, Any, Any]],
        iterations: int,
        warmup: int = 10
    ) -> list[float]:
        """Time an async operation multiple times."""
        # Warmup
        for _ in range(warmup):
            await operation()

        # Actual timing
        times = []
        for _ in range(iterations):
            start = time.perf_counter()
            await operation()
            end = time.perf_counter()
            times.append((end - start) * 1000)  # Convert to milliseconds

        return times

    def record_result(self, name: str, orm: str, times: list[float]):
        """Record benchmark results."""
        result = BenchmarkResult(
            name=name,
            orm=orm,
            mean_ms=statistics.mean(times),
            std_ms=statistics.stdev(times) if len(times) > 1 else 0,
            min_ms=min(times),
            max_ms=max(times),
            iterations=len(times),
        )
        self.results.append(result)
        return result


# ============= Benchmark Functions =============

class SQLAlchemyBenchmarks:
    """SQLAlchemy benchmark implementations."""

    def __init__(self, runner: BenchmarkRunner):
        self.runner = runner
        self._counter = 0

    def _unique_id(self) -> str:
        self._counter += 1
        return f"{uuid.uuid4().hex[:8]}_{self._counter}"

    async def single_insert(self) -> SAAuthor:
        """Insert a single author."""
        async with self.runner.sa_session() as session:
            author = SAAuthor(
                name=f"Author {self._unique_id()}",
                email=f"author_{self._unique_id()}@example.com",
                bio="A prolific writer"
            )
            session.add(author)
            await session.flush()
            return author

    async def bulk_insert(self, count: int = 100) -> list[SAAuthor]:
        """Insert multiple authors."""
        async with self.runner.sa_session() as session:
            authors = [
                SAAuthor(
                    name=f"Bulk Author {self._unique_id()}",
                    email=f"bulk_{self._unique_id()}@example.com",
                    bio=f"Bio for bulk author"
                )
                for _ in range(count)
            ]
            session.add_all(authors)
            await session.flush()
            return authors

    async def select_by_id(self, author_id: int) -> SAAuthor | None:
        """Select a single author by ID."""
        async with self.runner.sa_session() as session:
            result = await session.get(SAAuthor, author_id)
            return result

    async def select_with_filter(self) -> list[SAAuthor]:
        """Select authors with a filter."""
        async with self.runner.sa_session() as session:
            result = await session.execute(
                select(SAAuthor).where(SAAuthor.name.like("Bulk%")).limit(50)
            )
            return list(result.scalars().all())

    async def select_all(self, limit: int = 100) -> list[SAAuthor]:
        """Select all authors with limit."""
        async with self.runner.sa_session() as session:
            result = await session.execute(select(SAAuthor).limit(limit))
            return list(result.scalars().all())

    async def select_with_join(self) -> list[SABook]:
        """Select books with author and publisher (eager loading)."""
        async with self.runner.sa_session() as session:
            result = await session.execute(
                select(SABook)
                .options(selectinload(SABook.author), selectinload(SABook.publisher))
                .limit(50)
            )
            return list(result.scalars().all())

    async def update_single(self, author_id: int) -> None:
        """Update a single author."""
        async with self.runner.sa_session() as session:
            await session.execute(
                update(SAAuthor)
                .where(SAAuthor.id == author_id)
                .values(bio="Updated bio content")
            )

    async def update_bulk(self) -> None:
        """Update multiple authors."""
        async with self.runner.sa_session() as session:
            await session.execute(
                update(SAAuthor)
                .where(SAAuthor.name.like("Bulk%"))
                .values(bio="Bulk updated bio")
            )

    async def delete_single(self, author_id: int) -> None:
        """Delete a single author."""
        async with self.runner.sa_session() as session:
            await session.execute(
                delete(SAAuthor).where(SAAuthor.id == author_id)
            )

    async def aggregate_count(self) -> int:
        """Count all authors."""
        async with self.runner.sa_session() as session:
            result = await session.execute(select(sa_func.count(SAAuthor.id)))
            return result.scalar()

    async def aggregate_with_group(self) -> list:
        """Aggregate books by publisher with count."""
        async with self.runner.sa_session() as session:
            result = await session.execute(
                select(SAPublisher.name, sa_func.count(SABook.id))
                .join(SABook, SAPublisher.id == SABook.publisher_id, isouter=True)
                .group_by(SAPublisher.id, SAPublisher.name)
            )
            return list(result.all())


class TortoiseBenchmarks:
    """Tortoise ORM benchmark implementations."""

    def __init__(self):
        self._counter = 0

    def _unique_id(self) -> str:
        self._counter += 1
        return f"{uuid.uuid4().hex[:8]}_{self._counter}"

    async def single_insert(self) -> TAuthor:
        """Insert a single author."""
        author = await TAuthor.create(
            name=f"Author {self._unique_id()}",
            email=f"author_{self._unique_id()}@example.com",
            bio="A prolific writer"
        )
        return author

    async def bulk_insert(self, count: int = 100) -> list[TAuthor]:
        """Insert multiple authors."""
        authors = [
            TAuthor(
                name=f"Bulk Author {self._unique_id()}",
                email=f"bulk_{self._unique_id()}@example.com",
                bio=f"Bio for bulk author"
            )
            for _ in range(count)
        ]
        await TAuthor.bulk_create(authors)
        return authors

    async def select_by_id(self, author_id: int) -> TAuthor | None:
        """Select a single author by ID."""
        return await TAuthor.get_or_none(id=author_id)

    async def select_with_filter(self) -> list[TAuthor]:
        """Select authors with a filter."""
        return await TAuthor.filter(name__startswith="Bulk").limit(50)

    async def select_all(self, limit: int = 100) -> list[TAuthor]:
        """Select all authors with limit."""
        return await TAuthor.all().limit(limit)

    async def select_with_join(self) -> list[TBook]:
        """Select books with author and publisher (prefetch)."""
        return await TBook.all().prefetch_related("author", "publisher").limit(50)

    async def update_single(self, author_id: int) -> None:
        """Update a single author."""
        await TAuthor.filter(id=author_id).update(bio="Updated bio content")

    async def update_bulk(self) -> None:
        """Update multiple authors."""
        await TAuthor.filter(name__startswith="Bulk").update(bio="Bulk updated bio")

    async def delete_single(self, author_id: int) -> None:
        """Delete a single author."""
        await TAuthor.filter(id=author_id).delete()

    async def aggregate_count(self) -> int:
        """Count all authors."""
        return await TAuthor.all().count()

    async def aggregate_with_group(self) -> list:
        """Aggregate books by publisher with count."""
        return await TPublisher.all().annotate(book_count=Count("books")).values("name", "book_count")


async def setup_test_data(runner: BenchmarkRunner, tortoise_bench: TortoiseBenchmarks, sa_bench: SQLAlchemyBenchmarks):
    """Create test data for both ORMs."""
    console.print("[yellow]Setting up test data...[/yellow]")

    # SQLAlchemy test data
    async with runner.sa_session() as session:
        # Create publishers
        publishers = [
            SAPublisher(name=f"Publisher {i}", country="USA", founded_year=1990 + i)
            for i in range(10)
        ]
        session.add_all(publishers)
        await session.flush()

        # Create authors
        authors = [
            SAAuthor(name=f"Test Author {i}", email=f"test_{i}@example.com", bio=f"Bio {i}")
            for i in range(100)
        ]
        session.add_all(authors)
        await session.flush()

        # Create books
        books = [
            SABook(
                title=f"Book {i}",
                isbn=f"ISBN-SA-{i:05d}",
                pages=100 + i,
                price=Decimal("19.99"),
                published_date=date(2020, 1, 1),
                author_id=authors[i % len(authors)].id,
                publisher_id=publishers[i % len(publishers)].id,
            )
            for i in range(200)
        ]
        session.add_all(books)
        await session.flush()

    # Tortoise test data
    t_publishers = [
        await TPublisher.create(name=f"Publisher {i}", country="USA", founded_year=1990 + i)
        for i in range(10)
    ]

    t_authors = [
        await TAuthor.create(name=f"Test Author {i}", email=f"test_t_{i}@example.com", bio=f"Bio {i}")
        for i in range(100)
    ]

    t_books = [
        TBook(
            title=f"Book {i}",
            isbn=f"ISBN-T-{i:05d}",
            pages=100 + i,
            price=Decimal("19.99"),
            published_date=date(2020, 1, 1),
            author=t_authors[i % len(t_authors)],
            publisher=t_publishers[i % len(t_publishers)],
        )
        for i in range(200)
    ]
    await TBook.bulk_create(t_books)

    console.print("[green]Test data created successfully![/green]")

    return {
        "sa_authors": authors,
        "sa_publishers": publishers,
        "t_authors": t_authors,
        "t_publishers": t_publishers,
    }


async def run_benchmarks():
    """Run all benchmarks."""
    runner = BenchmarkRunner()
    iterations = BENCHMARK_SETTINGS["num_iterations"]
    warmup = BENCHMARK_SETTINGS["warmup_iterations"]

    console.print("\n[bold cyan]═══════════════════════════════════════════════════════════════[/bold cyan]")
    console.print("[bold cyan]     ORM Benchmark: TortoiseORM vs SQLAlchemy (asyncpg)         [/bold cyan]")
    console.print("[bold cyan]═══════════════════════════════════════════════════════════════[/bold cyan]\n")

    try:
        # Setup
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
        ) as progress:
            task = progress.add_task("Initializing SQLAlchemy...", total=None)
            await runner.setup_sqlalchemy()

            progress.update(task, description="Initializing Tortoise ORM...")
            await runner.setup_tortoise()

            progress.update(task, description="Setup complete!")

        sa_bench = SQLAlchemyBenchmarks(runner)
        t_bench = TortoiseBenchmarks()

        # Setup test data
        test_data = await setup_test_data(runner, t_bench, sa_bench)

        console.print(f"\n[cyan]Running benchmarks ({iterations} iterations each, {warmup} warmup)...[/cyan]\n")

        # ===== Single Insert Benchmark =====
        console.print("[yellow]1. Single Insert[/yellow]")

        times = await runner.time_operation(sa_bench.single_insert, iterations, warmup)
        runner.record_result("Single Insert", "SQLAlchemy", times)

        times = await runner.time_operation(t_bench.single_insert, iterations, warmup)
        runner.record_result("Single Insert", "TortoiseORM", times)

        # ===== Bulk Insert Benchmark =====
        console.print("[yellow]2. Bulk Insert (100 records)[/yellow]")

        times = await runner.time_operation(lambda: sa_bench.bulk_insert(100), iterations // 5, warmup // 5)
        runner.record_result("Bulk Insert (100)", "SQLAlchemy", times)

        times = await runner.time_operation(lambda: t_bench.bulk_insert(100), iterations // 5, warmup // 5)
        runner.record_result("Bulk Insert (100)", "TortoiseORM", times)

        # ===== Select by ID Benchmark =====
        console.print("[yellow]3. Select by ID[/yellow]")

        sa_author_id = test_data["sa_authors"][0].id
        t_author_id = test_data["t_authors"][0].id

        times = await runner.time_operation(lambda: sa_bench.select_by_id(sa_author_id), iterations, warmup)
        runner.record_result("Select by ID", "SQLAlchemy", times)

        times = await runner.time_operation(lambda: t_bench.select_by_id(t_author_id), iterations, warmup)
        runner.record_result("Select by ID", "TortoiseORM", times)

        # ===== Select with Filter Benchmark =====
        console.print("[yellow]4. Select with Filter[/yellow]")

        times = await runner.time_operation(sa_bench.select_with_filter, iterations, warmup)
        runner.record_result("Select with Filter", "SQLAlchemy", times)

        times = await runner.time_operation(t_bench.select_with_filter, iterations, warmup)
        runner.record_result("Select with Filter", "TortoiseORM", times)

        # ===== Select All (Limit 100) Benchmark =====
        console.print("[yellow]5. Select All (Limit 100)[/yellow]")

        times = await runner.time_operation(lambda: sa_bench.select_all(100), iterations, warmup)
        runner.record_result("Select All (100)", "SQLAlchemy", times)

        times = await runner.time_operation(lambda: t_bench.select_all(100), iterations, warmup)
        runner.record_result("Select All (100)", "TortoiseORM", times)

        # ===== Select with Join Benchmark =====
        console.print("[yellow]6. Select with JOIN (Eager Loading)[/yellow]")

        times = await runner.time_operation(sa_bench.select_with_join, iterations, warmup)
        runner.record_result("Select with JOIN", "SQLAlchemy", times)

        times = await runner.time_operation(t_bench.select_with_join, iterations, warmup)
        runner.record_result("Select with JOIN", "TortoiseORM", times)

        # ===== Update Single Benchmark =====
        console.print("[yellow]7. Update Single Record[/yellow]")

        times = await runner.time_operation(lambda: sa_bench.update_single(sa_author_id), iterations, warmup)
        runner.record_result("Update Single", "SQLAlchemy", times)

        times = await runner.time_operation(lambda: t_bench.update_single(t_author_id), iterations, warmup)
        runner.record_result("Update Single", "TortoiseORM", times)

        # ===== Update Bulk Benchmark =====
        console.print("[yellow]8. Update Bulk[/yellow]")

        times = await runner.time_operation(sa_bench.update_bulk, iterations // 2, warmup // 2)
        runner.record_result("Update Bulk", "SQLAlchemy", times)

        times = await runner.time_operation(t_bench.update_bulk, iterations // 2, warmup // 2)
        runner.record_result("Update Bulk", "TortoiseORM", times)

        # ===== Aggregate Count Benchmark =====
        console.print("[yellow]9. Aggregate Count[/yellow]")

        times = await runner.time_operation(sa_bench.aggregate_count, iterations, warmup)
        runner.record_result("Aggregate Count", "SQLAlchemy", times)

        times = await runner.time_operation(t_bench.aggregate_count, iterations, warmup)
        runner.record_result("Aggregate Count", "TortoiseORM", times)

        # ===== Aggregate with Group Benchmark =====
        console.print("[yellow]10. Aggregate with GROUP BY[/yellow]")

        times = await runner.time_operation(sa_bench.aggregate_with_group, iterations, warmup)
        runner.record_result("Aggregate GROUP BY", "SQLAlchemy", times)

        times = await runner.time_operation(t_bench.aggregate_with_group, iterations, warmup)
        runner.record_result("Aggregate GROUP BY", "TortoiseORM", times)

        # Print results
        print_results(runner.results)

    finally:
        # Cleanup
        console.print("\n[yellow]Cleaning up...[/yellow]")
        await runner.cleanup_tortoise()
        await runner.cleanup_sqlalchemy()
        console.print("[green]Cleanup complete![/green]")


def print_results(results: list[BenchmarkResult]):
    """Print benchmark results in a nice table."""
    console.print("\n")
    console.print("[bold cyan]═══════════════════════════════════════════════════════════════[/bold cyan]")
    console.print("[bold cyan]                      BENCHMARK RESULTS                         [/bold cyan]")
    console.print("[bold cyan]═══════════════════════════════════════════════════════════════[/bold cyan]\n")

    # Group results by benchmark name
    benchmarks = {}
    for r in results:
        if r.name not in benchmarks:
            benchmarks[r.name] = {}
        benchmarks[r.name][r.orm] = r

    # Create comparison table
    table = Table(title="Performance Comparison (times in milliseconds)")
    table.add_column("Benchmark", style="cyan", no_wrap=True)
    table.add_column("SQLAlchemy Mean", justify="right", style="green")
    table.add_column("TortoiseORM Mean", justify="right", style="yellow")
    table.add_column("Winner", justify="center", style="bold")
    table.add_column("Difference", justify="right")

    for name, orms in benchmarks.items():
        sa_result = orms.get("SQLAlchemy")
        t_result = orms.get("TortoiseORM")

        if sa_result and t_result:
            sa_mean = f"{sa_result.mean_ms:.3f}"
            t_mean = f"{t_result.mean_ms:.3f}"

            if sa_result.mean_ms < t_result.mean_ms:
                winner = "[green]SQLAlchemy[/green]"
                diff_pct = ((t_result.mean_ms - sa_result.mean_ms) / t_result.mean_ms) * 100
                diff = f"[green]{diff_pct:.1f}% faster[/green]"
            else:
                winner = "[yellow]TortoiseORM[/yellow]"
                diff_pct = ((sa_result.mean_ms - t_result.mean_ms) / sa_result.mean_ms) * 100
                diff = f"[yellow]{diff_pct:.1f}% faster[/yellow]"

            table.add_row(name, sa_mean, t_mean, winner, diff)

    console.print(table)

    # Detailed results
    console.print("\n[bold]Detailed Results:[/bold]\n")

    detail_table = Table(title="Detailed Statistics")
    detail_table.add_column("Benchmark", style="cyan")
    detail_table.add_column("ORM", style="white")
    detail_table.add_column("Mean (ms)", justify="right")
    detail_table.add_column("Std Dev", justify="right")
    detail_table.add_column("Min (ms)", justify="right")
    detail_table.add_column("Max (ms)", justify="right")
    detail_table.add_column("Iterations", justify="right")

    for r in results:
        detail_table.add_row(
            r.name,
            r.orm,
            f"{r.mean_ms:.3f}",
            f"{r.std_ms:.3f}",
            f"{r.min_ms:.3f}",
            f"{r.max_ms:.3f}",
            str(r.iterations),
        )

    console.print(detail_table)

    # Summary
    console.print("\n[bold cyan]Summary:[/bold cyan]")
    sa_wins = sum(1 for name, orms in benchmarks.items()
                  if orms.get("SQLAlchemy") and orms.get("TortoiseORM")
                  and orms["SQLAlchemy"].mean_ms < orms["TortoiseORM"].mean_ms)
    t_wins = len(benchmarks) - sa_wins

    console.print(f"  SQLAlchemy wins: [green]{sa_wins}[/green] benchmarks")
    console.print(f"  TortoiseORM wins: [yellow]{t_wins}[/yellow] benchmarks")


if __name__ == "__main__":
    asyncio.run(run_benchmarks())