main.py

#!/usr/bin/env python3
"""
ORM Benchmarks: TortoiseORM vs SQLAlchemy

Entry point for running ORM benchmarks comparing TortoiseORM and SQLAlchemy,
both using the asyncpg driver for async PostgreSQL operations.

Usage:
    python main.py              # Run benchmarks with default settings
    python main.py --help       # Show help
"""

import argparse
import asyncio
import os
import sys


def main():
    parser = argparse.ArgumentParser(
        description="Benchmark TortoiseORM vs SQLAlchemy with asyncpg"
    )
    parser.add_argument(
        "--db-host",
        default=os.getenv("DB_HOST", "localhost"),
        help="Database host (default: localhost or DB_HOST env var)",
    )
    parser.add_argument(
        "--db-port",
        type=int,
        default=int(os.getenv("DB_PORT", "5432")),
        help="Database port (default: 5432 or DB_PORT env var)",
    )
    parser.add_argument(
        "--db-user",
        default=os.getenv("DB_USER", "postgres"),
        help="Database user (default: postgres or DB_USER env var)",
    )
    parser.add_argument(
        "--db-password",
        default=os.getenv("DB_PASSWORD", "postgres"),
        help="Database password (default: postgres or DB_PASSWORD env var)",
    )
    parser.add_argument(
        "--db-name",
        default=os.getenv("DB_NAME", "orm_benchmark"),
        help="Database name (default: orm_benchmark or DB_NAME env var)",
    )
    parser.add_argument(
        "--iterations",
        type=int,
        default=100,
        help="Number of iterations per benchmark (default: 100)",
    )
    parser.add_argument(
        "--warmup",
        type=int,
        default=10,
        help="Number of warmup iterations (default: 10)",
    )

    args = parser.parse_args()

    # Set environment variables for config module
    os.environ["DB_HOST"] = args.db_host
    os.environ["DB_PORT"] = str(args.db_port)
    os.environ["DB_USER"] = args.db_user
    os.environ["DB_PASSWORD"] = args.db_password
    os.environ["DB_NAME"] = args.db_name

    # Update benchmark settings
    from config import BENCHMARK_SETTINGS
    BENCHMARK_SETTINGS["num_iterations"] = args.iterations
    BENCHMARK_SETTINGS["warmup_iterations"] = args.warmup

    # Run benchmarks
    from benchmarks import run_benchmarks

    try:
        asyncio.run(run_benchmarks())
    except KeyboardInterrupt:
        print("\nBenchmark interrupted.")
        sys.exit(1)
    except Exception as e:
        print(f"\nError running benchmarks: {e}")
        print("\nMake sure PostgreSQL is running and the database exists.")
        print("You can create the database with:")
        print(f"  createdb {args.db_name}")
        print("\nOr set the connection details via environment variables or command line args.")
        sys.exit(1)


if __name__ == "__main__":
    main()