psycopg3 vs psycopg2 description benchmark

benchmark.py

# Run via: uv run benchmark.py
# /// script
# dependencies = [
#    "psycopg2>=2.9.11",
#    "psycopg[c]>=3.3.2",
#    "py-spy>=0.4.1",
# ]
# ///


import psycopg
import psycopg2
from psycopg import sql
import subprocess
import sys
import time
import os


def get_db_user():
    """Get the database user, preferring the original user if running with sudo"""
    return os.environ.get("SUDO_USER") or os.environ.get("USER")


def setup_database():
    # Connect to the default postgres database to create our database
    db_user = get_db_user()
    conn_params = {
        "host": "localhost",
        "dbname": "postgres",
        "user": db_user,
    }

    # Create database if it doesn't exist
    with psycopg.connect(**conn_params, autocommit=True) as conn:
        with conn.cursor() as cur:
            # Check if database exists
            cur.execute(
                "SELECT 1 FROM pg_database WHERE datname = %s",
                ("psycopg_benchmark",)
            )
            if not cur.fetchone():
                cur.execute(
                    sql.SQL("CREATE DATABASE {}").format(
                        sql.Identifier("psycopg_benchmark")
                    )
                )
                print("Database 'psycopg_benchmark' created.")
            else:
                print("Database 'psycopg_benchmark' already exists.")

    # Connect to the psycopg_benchmark database
    conn_params["dbname"] = "psycopg_benchmark"
    with psycopg.connect(**conn_params) as conn:
        with conn.cursor() as cur:
            # Create table if it doesn't exist
            cur.execute("""
                CREATE TABLE IF NOT EXISTS benchmark_table (
                    id SERIAL PRIMARY KEY,
                    column_1 BIGINT NOT NULL,
                    column_2 BOOLEAN NOT NULL,
                    column_3 TIMESTAMP WITHOUT TIME ZONE NOT NULL,
                    column_4 TIMESTAMP WITHOUT TIME ZONE NOT NULL,
                    column_5 TIMESTAMP WITHOUT TIME ZONE,
                    column_6 VARCHAR(3),
                    column_7 VARCHAR,
                    column_8 BIGINT,
                    column_9 BOOLEAN,
                    column_10 JSON,
                    column_11 JSON,
                    column_12 VARCHAR,
                    column_13 BOOLEAN,
                    column_14 VARCHAR,
                    column_15 BIGINT,
                    column_16 TIMESTAMP WITHOUT TIME ZONE,
                    column_17 VARCHAR(2),
                    column_18 VARCHAR(2),
                    column_19 SMALLINT,
                    column_20 VARCHAR(128),
                    column_21 VARCHAR(128),
                    column_22 BOOLEAN DEFAULT FALSE NOT NULL
                )
            """)
            conn.commit()
            print("Table 'benchmark_table' created or already exists.")


def benchmark_description(conn, iterations=10000):
    """Benchmark cursor.description access"""
    cur = conn.cursor()
    cur.execute("SELECT * FROM benchmark_table LIMIT 1")

    # Access description many times
    for _ in range(iterations):
        desc = cur.description

    cur.close()


def run_psycopg3():
    """Run benchmark with psycopg3"""
    db_user = get_db_user()
    conn_params = {"host": "localhost", "dbname": "psycopg_benchmark", "user": db_user}
    start = time.perf_counter()
    with psycopg.connect(**conn_params) as conn:
        benchmark_description(conn, iterations=100000)
    elapsed = time.perf_counter() - start
    print(f"psycopg3: Accessed cursor.description 100000 times in {elapsed:.9f} seconds")


def run_psycopg2():
    """Run benchmark with psycopg2"""
    db_user = get_db_user()
    conn_params = {"host": "localhost", "dbname": "psycopg_benchmark", "user": db_user}
    start = time.perf_counter()
    conn = psycopg2.connect(**conn_params)
    try:
        benchmark_description(conn, iterations=100000)
    finally:
        conn.close()
    elapsed = time.perf_counter() - start
    print(f"psycopg2: Accessed cursor.description 100000 times in {elapsed:.9f} seconds")


def main():
    db_user = get_db_user()
    print(f"Using database user: {db_user}")
    if os.environ.get("SUDO_USER"):
        print("(Running with elevated permissions, connecting as original user)")

    print("\nSetting up database...")
    setup_database()

    print("\n" + "="*60)
    print("Running benchmarks with py-spy profiling")
    print("="*60 + "\n")

    # Profile psycopg3
    print("Profiling psycopg3 cursor.description...")
    result = subprocess.run([
        "py-spy", "record",
        "-o", "psycopg3_description_profile.svg",
        "--", sys.executable, "-c",
        "from main import run_psycopg3; run_psycopg3()"
    ], capture_output=True, text=True)

    if result.returncode == 0:
        print("✓ psycopg3 profile saved to psycopg3_description_profile.svg")
    else:
        print("⚠ Warning: py-spy failed (may need elevated permissions)")
        print("Running psycopg3 benchmark without profiling...")
        run_psycopg3()

    # Profile psycopg2
    print("\nProfiling psycopg2 cursor.description...")
    result = subprocess.run([
        "py-spy", "record",
        "-o", "psycopg2_description_profile.svg",
        "--", sys.executable, "-c",
        "from main import run_psycopg2; run_psycopg2()"
    ], capture_output=True, text=True)

    if result.returncode == 0:
        print("✓ psycopg2 profile saved to psycopg2_description_profile.svg")
    else:
        print("⚠ Warning: py-spy failed (may need elevated permissions)")
        print("Running psycopg2 benchmark without profiling...")
        run_psycopg2()

    print("\nBenchmark complete!")
    print("\nTo enable profiling, run with: sudo python main.py")
    print("Or run py-spy manually: sudo py-spy record -o profile.svg -- python -c 'from main import run_psycopg3; run_psycopg3()'")


if __name__ == "__main__":
    main()