iceaxe 0.8.3__cp313-cp313-macosx_11_0_arm64.whl

iceaxe/__init__.py

@@ -0,0 +1,20 @@
+from .alias_values import alias as alias
+from .base import (
+    IndexConstraint as IndexConstraint,
+    TableBase as TableBase,
+    UniqueConstraint as UniqueConstraint,
+)
+from .field import Field as Field
+from .functions import func as func
+from .postgres import PostgresDateTime as PostgresDateTime, PostgresTime as PostgresTime
+from .queries import (
+    QueryBuilder as QueryBuilder,
+    and_ as and_,
+    delete as delete,
+    or_ as or_,
+    select as select,
+    update as update,
+)
+from .queries_str import sql as sql
+from .session import DBConnection as DBConnection
+from .typing import column as column

iceaxe/__tests__/benchmarks/test_bulk_insert.py

@@ -0,0 +1,45 @@
+import time
+from typing import Sequence
+
+import pytest
+
+from iceaxe.__tests__.conf_models import UserDemo
+from iceaxe.logging import CONSOLE, LOGGER
+from iceaxe.session import DBConnection
+
+
+def generate_test_users(count: int) -> Sequence[UserDemo]:
+    """
+    Generate a sequence of test users for bulk insertion.
+
+    :param count: Number of users to generate
+    :return: Sequence of UserDemo instances
+    """
+    return [
+        UserDemo(name=f"User {i}", email=f"user{i}@example.com") for i in range(count)
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration_tests
+async def test_bulk_insert_performance(db_connection: DBConnection):
+    """
+    Test the performance of bulk inserting 500k records.
+    """
+    NUM_USERS = 500_000
+    users = generate_test_users(NUM_USERS)
+    LOGGER.info(f"Generated {NUM_USERS} test users")
+
+    start_time = time.time()
+
+    await db_connection.insert(users)
+
+    total_time = time.time() - start_time
+    records_per_second = NUM_USERS / total_time
+
+    CONSOLE.print("\nBulk Insert Performance:")
+    CONSOLE.print(f"Total time: {total_time:.2f} seconds")
+    CONSOLE.print(f"Records per second: {records_per_second:.2f}")
+
+    result = await db_connection.conn.fetchval("SELECT COUNT(*) FROM userdemo")
+    assert result == NUM_USERS

iceaxe/__tests__/benchmarks/test_select.py

@@ -0,0 +1,114 @@
+from enum import Enum
+from time import monotonic_ns
+from typing import Any
+
+import asyncpg
+import pytest
+
+from iceaxe.__tests__.conf_models import UserDemo, run_profile
+from iceaxe.logging import CONSOLE, LOGGER
+from iceaxe.queries import QueryBuilder
+from iceaxe.session import DBConnection
+
+
+class FetchType(Enum):
+    ID = "id"
+    OBJ = "obj"
+
+
+async def insert_users(conn: asyncpg.Connection, num_users: int):
+    users = [(f"User {i}", f"user{i}@example.com") for i in range(num_users)]
+    await conn.executemany("INSERT INTO userdemo (name, email) VALUES ($1, $2)", users)
+
+
+async def fetch_users_raw(conn: asyncpg.Connection, fetch_type: FetchType) -> list[Any]:
+    if fetch_type == FetchType.OBJ:
+        return await conn.fetch("SELECT * FROM userdemo")  # type: ignore
+    elif fetch_type == FetchType.ID:
+        return await conn.fetch("SELECT id FROM userdemo")  # type: ignore
+    else:
+        raise ValueError(f"Invalid run profile: {fetch_type}")
+
+
+def build_iceaxe_query(fetch_type: FetchType):
+    if fetch_type == FetchType.OBJ:
+        return QueryBuilder().select(UserDemo)
+    elif fetch_type == FetchType.ID:
+        return QueryBuilder().select(UserDemo.id)
+    else:
+        raise ValueError(f"Invalid run profile: {fetch_type}")
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration_tests
+@pytest.mark.parametrize(
+    "fetch_type, allowed_overhead",
+    [
+        (FetchType.ID, 10),
+        (FetchType.OBJ, 800),
+    ],
+)
+async def test_benchmark(
+    db_connection: DBConnection, request, fetch_type: FetchType, allowed_overhead: float
+):
+    num_users = 500_000
+    num_loops = 100
+
+    # Insert users using raw asyncpg
+    await insert_users(db_connection.conn, num_users)
+
+    # Benchmark raw asyncpg query
+    start_time = monotonic_ns()
+    raw_results: list[Any] = []
+    for _ in range(num_loops):
+        raw_results = await fetch_users_raw(db_connection.conn, fetch_type)
+    raw_time = monotonic_ns() - start_time
+    raw_time_seconds = raw_time / 1e9
+    raw_time_per_query = (raw_time / num_loops) / 1e9
+
+    LOGGER.info(
+        f"Raw asyncpg query time: {raw_time_per_query:.4f} (total: {raw_time_seconds:.4f}) seconds"
+    )
+    CONSOLE.print(
+        f"Raw asyncpg query time: {raw_time_per_query:.4f} (total: {raw_time_seconds:.4f}) seconds"
+    )
+
+    # Benchmark DBConnection.exec query
+    start_time = monotonic_ns()
+    query = build_iceaxe_query(fetch_type)
+    db_results: list[UserDemo] | list[int] = []
+    for _ in range(num_loops):
+        db_results = await db_connection.exec(query)
+    db_time = monotonic_ns() - start_time
+    db_time_seconds = db_time / 1e9
+    db_time_per_query = (db_time / num_loops) / 1e9
+
+    LOGGER.info(
+        f"DBConnection.exec query time: {db_time_per_query:.4f} (total: {db_time_seconds:.4f}) seconds"
+    )
+    CONSOLE.print(
+        f"DBConnection.exec query time: {db_time_per_query:.4f} (total: {db_time_seconds:.4f}) seconds"
+    )
+
+    # Slower than the raw run since we need to run the performance instrumentation
+    if False:
+        with run_profile(request):
+            # Right now we don't cache results so we can run multiple times to get a better measure of samples
+            for _ in range(num_loops):
+                query = build_iceaxe_query(fetch_type)
+                db_results = await db_connection.exec(query)
+
+    # Compare results
+    assert len(raw_results) == len(db_results) == num_users, "Result count mismatch"
+
+    # Calculate performance difference
+    performance_diff = (db_time - raw_time) / raw_time * 100
+    LOGGER.info(f"Performance difference: {performance_diff:.2f}%")
+    CONSOLE.print(f"Performance difference: {performance_diff:.2f}%")
+
+    # Assert that DBConnection.exec is at most X% slower than raw query
+    assert performance_diff <= allowed_overhead, (
+        f"DBConnection.exec is {performance_diff:.2f}% slower than raw query, which exceeds the {allowed_overhead}% threshold"
+    )
+
+    LOGGER.info("Benchmark completed successfully.")

iceaxe/__tests__/conf_models.py

@@ -0,0 +1,133 @@
+from contextlib import contextmanager
+from enum import StrEnum
+from pathlib import Path
+from typing import Any
+
+from pyinstrument import Profiler
+
+from iceaxe.base import Field, TableBase, UniqueConstraint
+
+
+class UserDemo(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    name: str
+    email: str
+
+
+class ArtifactDemo(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    title: str
+    user_id: int = Field(foreign_key="userdemo.id")
+
+
+class ComplexDemo(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    string_list: list[str]
+    json_data: dict[str, str] = Field(is_json=True)
+
+
+class Employee(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    email: str = Field(unique=True)
+    first_name: str
+    last_name: str
+    department: str
+    salary: float
+
+
+class Department(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    name: str = Field(unique=True)
+    budget: float
+    location: str
+
+
+class ProjectAssignment(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    employee_id: int = Field(foreign_key="employee.id")
+    project_name: str
+    role: str
+    start_date: str
+
+
+class EmployeeStatus(StrEnum):
+    ACTIVE = "active"
+    INACTIVE = "inactive"
+    ON_LEAVE = "on_leave"
+
+
+class EmployeeMetadata(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    employee_id: int = Field(foreign_key="employee.id")
+    status: EmployeeStatus
+    tags: list[str] = Field(is_json=True)
+    additional_info: dict[str, Any] = Field(is_json=True)
+
+
+class FunctionDemoModel(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    balance: float
+    created_at: str
+    birth_date: str
+    start_date: str
+    end_date: str
+    year: int
+    month: int
+    day: int
+    hour: int
+    minute: int
+    second: int
+    years: int
+    months: int
+    days: int
+    weeks: int
+    hours: int
+    minutes: int
+    seconds: int
+    name: str
+    balance_str: str
+    timestamp_str: str
+
+
+class DemoModelA(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    name: str
+    description: str
+    code: str = Field(unique=True)
+
+
+class DemoModelB(TableBase):
+    id: int = Field(primary_key=True, default=None)
+    name: str
+    category: str
+    code: str = Field(unique=True)
+
+
+class JsonDemo(TableBase):
+    """
+    Model for testing JSON field updates.
+    """
+
+    id: int | None = Field(primary_key=True, default=None)
+    settings: dict[Any, Any] = Field(is_json=True)
+    metadata: dict[Any, Any] | None = Field(is_json=True)
+    unique_val: str
+
+    table_args = [UniqueConstraint(columns=["unique_val"])]
+
+
+@contextmanager
+def run_profile(request):
+    TESTS_ROOT = Path.cwd()
+    PROFILE_ROOT = TESTS_ROOT / ".profiles"
+
+    # Turn profiling on
+    profiler = Profiler()
+    profiler.start()
+
+    yield  # Run test
+
+    profiler.stop()
+    PROFILE_ROOT.mkdir(exist_ok=True)
+    results_file = PROFILE_ROOT / f"{request.node.name}.html"
+    profiler.write_html(results_file)

iceaxe/__tests__/conftest.py

@@ -0,0 +1,204 @@
+import logging
+
+import asyncpg
+import pytest
+import pytest_asyncio
+
+from iceaxe.__tests__ import docker_helpers
+from iceaxe.base import DBModelMetaclass
+from iceaxe.session import DBConnection
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@pytest.fixture(scope="session")
+def docker_postgres():
+    """
+    Fixture that creates a PostgreSQL container using the Python Docker API.
+    This allows running individual tests without needing Docker Compose.
+    """
+    # Create and start a PostgreSQL container
+    postgres_container = docker_helpers.PostgresContainer()
+
+    # Start the container and yield connection details
+    connection_info = postgres_container.start()
+    yield connection_info
+
+    # Cleanup: stop the container
+    postgres_container.stop()
+
+
+@pytest_asyncio.fixture
+async def db_connection(docker_postgres):
+    """
+    Create a database connection using the PostgreSQL container.
+    """
+    conn = DBConnection(
+        await asyncpg.connect(
+            host=docker_postgres["host"],
+            port=docker_postgres["port"],
+            user=docker_postgres["user"],
+            password=docker_postgres["password"],
+            database=docker_postgres["database"],
+        )
+    )
+
+    # Drop all tables first to ensure clean state
+    known_tables = [
+        "artifactdemo",
+        "userdemo",
+        "complexdemo",
+        "article",
+        "employee",
+        "department",
+        "projectassignment",
+        "employeemetadata",
+        "functiondemomodel",
+        "demomodela",
+        "demomodelb",
+        "jsondemo",
+        "complextypedemo",
+    ]
+    known_types = ["statusenum", "employeestatus"]
+
+    for table in known_tables:
+        await conn.conn.execute(f"DROP TABLE IF EXISTS {table} CASCADE", timeout=30.0)
+
+    for known_type in known_types:
+        await conn.conn.execute(
+            f"DROP TYPE IF EXISTS {known_type} CASCADE", timeout=30.0
+        )
+
+    # Create tables
+    await conn.conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS userdemo (
+            id SERIAL PRIMARY KEY,
+            name TEXT,
+            email TEXT
+        )
+    """,
+        timeout=30.0,
+    )
+
+    await conn.conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS artifactdemo (
+            id SERIAL PRIMARY KEY,
+            title TEXT,
+            user_id INT REFERENCES userdemo(id)
+        )
+    """,
+        timeout=30.0,
+    )
+
+    await conn.conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS complexdemo (
+            id SERIAL PRIMARY KEY,
+            string_list TEXT[],
+            json_data JSON
+        )
+    """,
+        timeout=30.0,
+    )
+
+    await conn.conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS article (
+            id SERIAL PRIMARY KEY,
+            title TEXT,
+            content TEXT,
+            summary TEXT
+        )
+    """,
+        timeout=30.0,
+    )
+
+    # Create each index separately to handle errors better
+    yield conn
+
+    # Drop all tables after tests
+    for table in known_tables:
+        await conn.conn.execute(f"DROP TABLE IF EXISTS {table} CASCADE", timeout=30.0)
+
+    # Drop all types after tests
+    for known_type in known_types:
+        await conn.conn.execute(
+            f"DROP TYPE IF EXISTS {known_type} CASCADE", timeout=30.0
+        )
+
+    await conn.conn.close()
+
+
+@pytest_asyncio.fixture()
+async def indexed_db_connection(db_connection: DBConnection):
+    await db_connection.conn.execute(
+        "CREATE INDEX IF NOT EXISTS article_title_tsv_idx ON article USING GIN (to_tsvector('english', title))",
+        timeout=30.0,
+    )
+    await db_connection.conn.execute(
+        "CREATE INDEX IF NOT EXISTS article_content_tsv_idx ON article USING GIN (to_tsvector('english', content))",
+        timeout=30.0,
+    )
+    await db_connection.conn.execute(
+        "CREATE INDEX IF NOT EXISTS article_summary_tsv_idx ON article USING GIN (to_tsvector('english', summary))",
+        timeout=30.0,
+    )
+
+    yield db_connection
+
+
+@pytest_asyncio.fixture(autouse=True)
+async def clear_table(db_connection):
+    # Clear all tables and reset sequences
+    await db_connection.conn.execute(
+        "TRUNCATE TABLE userdemo, article RESTART IDENTITY CASCADE", timeout=30.0
+    )
+
+
+@pytest_asyncio.fixture
+async def clear_all_database_objects(db_connection: DBConnection):
+    """
+    Clear all database objects.
+    """
+    # Step 1: Drop all tables in the public schema
+    await db_connection.conn.execute(
+        """
+        DO $$ DECLARE
+            r RECORD;
+        BEGIN
+            FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public') LOOP
+                EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';
+            END LOOP;
+        END $$;
+    """,
+        timeout=30.0,
+    )
+
+    # Step 2: Drop all custom types in the public schema
+    await db_connection.conn.execute(
+        """
+        DO $$ DECLARE
+            r RECORD;
+        BEGIN
+            FOR r IN (SELECT typname FROM pg_type WHERE typtype = 'e' AND typnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')) LOOP
+                EXECUTE 'DROP TYPE IF EXISTS ' || quote_ident(r.typname) || ' CASCADE';
+            END LOOP;
+        END $$;
+    """,
+        timeout=30.0,
+    )
+
+
+@pytest.fixture
+def clear_registry():
+    current_registry = DBModelMetaclass._registry
+    DBModelMetaclass._registry = []
+
+    try:
+        yield
+    finally:
+        DBModelMetaclass._registry = current_registry