iceaxe 0.7.0.dev3__tar.gz → 0.7.1__tar.gz

{iceaxe-0.7.0.dev3/iceaxe.egg-info → iceaxe-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iceaxe
-Version: 0.7.0.dev3
+Version: 0.7.1
 Summary: A modern, fast ORM for Python.
 Author-email: Pierce Freeman <pierce@freeman.vc>
 Requires-Python: >=3.11

{iceaxe-0.7.0.dev3 → iceaxe-0.7.1}/iceaxe/__tests__/conftest.py

@@ -1,84 +1,151 @@
+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():
+async def db_connection(docker_postgres):
+    """
+    Create a database connection using the PostgreSQL container.
+    """
     conn = DBConnection(
         await asyncpg.connect(
-            host="localhost",
-            port=5438,
-            user="iceaxe",
-            password="mysecretpassword",
-            database="iceaxe_test_db",
+            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
-    await conn.conn.execute("DROP TABLE IF EXISTS artifactdemo CASCADE")
-    await conn.conn.execute("DROP TABLE IF EXISTS userdemo CASCADE")
-    await conn.conn.execute("DROP TABLE IF EXISTS complexdemo CASCADE")
-    await conn.conn.execute("DROP TABLE IF EXISTS article CASCADE")
+    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("""
+    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("""
+    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("""
+    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("""
+    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
-    await conn.conn.execute("DROP TABLE IF EXISTS artifactdemo CASCADE")
-    await conn.conn.execute("DROP TABLE IF EXISTS userdemo CASCADE")
-    await conn.conn.execute("DROP TABLE IF EXISTS complexdemo CASCADE")
-    await conn.conn.execute("DROP TABLE IF EXISTS article CASCADE")
+    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))"
+        "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))"
+        "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))"
+        "CREATE INDEX IF NOT EXISTS article_summary_tsv_idx ON article USING GIN (to_tsvector('english', summary))",
+        timeout=30.0,
     )
 
     yield db_connection
@@ -88,7 +155,7 @@ async def indexed_db_connection(db_connection: DBConnection):
 async def clear_table(db_connection):
     # Clear all tables and reset sequences
     await db_connection.conn.execute(
-        "TRUNCATE TABLE userdemo, article RESTART IDENTITY CASCADE"
+        "TRUNCATE TABLE userdemo, article RESTART IDENTITY CASCADE", timeout=30.0
     )
 
 
@@ -107,7 +174,8 @@ async def clear_all_database_objects(db_connection: DBConnection):
                 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
@@ -120,7 +188,8 @@ async def clear_all_database_objects(db_connection: DBConnection):
                 EXECUTE 'DROP TYPE IF EXISTS ' || quote_ident(r.typname) || ' CASCADE';
             END LOOP;
         END $$;
-    """
+    """,
+        timeout=30.0,
     )
 
 

iceaxe-0.7.1/iceaxe/__tests__/docker_helpers.py

@@ -0,0 +1,208 @@
+"""
+Docker helper utilities for testing.
+
+This module provides classes and functions to manage Docker containers for testing,
+particularly focusing on PostgreSQL database containers.
+"""
+
+import logging
+import socket
+import time
+import uuid
+from typing import Any, Dict, Optional, cast
+
+import docker
+from docker.errors import APIError
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+
+def get_free_port() -> int:
+    """Find a free port on the host machine."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+class PostgresContainer:
+    """
+    A class that manages a PostgreSQL Docker container for testing.
+
+    This class handles the lifecycle of a PostgreSQL container, including:
+    - Starting the container with appropriate configuration
+    - Finding available ports
+    - Waiting for the container to be ready
+    - Providing connection information
+    - Cleaning up after tests
+    """
+
+    def __init__(
+        self,
+        pg_user: str = "iceaxe",
+        pg_password: str = "mysecretpassword",
+        pg_db: str = "iceaxe_test_db",
+        postgres_version: str = "16",
+    ):
+        self.pg_user = pg_user
+        self.pg_password = pg_password
+        self.pg_db = pg_db
+        self.postgres_version = postgres_version
+        self.port = get_free_port()
+        self.container: Optional[Any] = None
+        self.client = docker.from_env()
+        self.container_name = f"iceaxe-postgres-test-{uuid.uuid4().hex[:8]}"
+
+    def start(self) -> Dict[str, Any]:
+        """
+        Start the PostgreSQL container.
+
+        Returns:
+            Dict[str, Any]: Connection information for the PostgreSQL container
+
+        Raises:
+            RuntimeError: If the container fails to start or become ready
+        """
+        logger.info(f"Starting PostgreSQL container on port {self.port}")
+
+        max_attempts = 3
+        attempt = 0
+
+        while attempt < max_attempts:
+            attempt += 1
+            try:
+                self.container = self._run_container(self.port)
+                break
+            except APIError as e:
+                if "port is already allocated" in str(e) and attempt < max_attempts:
+                    logger.warning(
+                        f"Port {self.port} is still in use. Trying with a new port (attempt {attempt}/{max_attempts})."
+                    )
+                    self.port = get_free_port()
+                else:
+                    raise RuntimeError(f"Failed to start PostgreSQL container: {e}")
+
+        # Wait for PostgreSQL to be ready
+        if not self._wait_for_container_ready():
+            self.stop()
+            raise RuntimeError("Failed to connect to PostgreSQL container")
+
+        return self.get_connection_info()
+
+    def _run_container(
+        self, port: int
+    ) -> Any:  # Type as Any since docker.models.containers.Container isn't imported
+        """
+        Run the Docker container with the specified port.
+
+        Args:
+            port: The port to map PostgreSQL to on the host
+
+        Returns:
+            The Docker container object
+        """
+        return self.client.containers.run(
+            f"postgres:{self.postgres_version}",
+            name=self.container_name,
+            detach=True,
+            environment={
+                "POSTGRES_USER": self.pg_user,
+                "POSTGRES_PASSWORD": self.pg_password,
+                "POSTGRES_DB": self.pg_db,
+                # Additional settings for faster startup in testing
+                "POSTGRES_HOST_AUTH_METHOD": "trust",
+            },
+            ports={"5432/tcp": port},
+            remove=True,  # Auto-remove container when stopped
+        )
+
+    def _wait_for_container_ready(self) -> bool:
+        """
+        Wait for the PostgreSQL container to be ready.
+
+        Returns:
+            bool: True if the container is ready, False otherwise
+        """
+        max_retries = 30
+        retry_interval = 1
+
+        for i in range(max_retries):
+            try:
+                if self.container is None:
+                    logger.warning("Container is None, cannot proceed")
+                    return False
+
+                # We've already checked that self.container is not None
+                container = cast(Any, self.container)
+                container.reload()  # Refresh container status
+                if container.status != "running":
+                    logger.warning(f"Container status: {container.status}")
+                    return False
+
+                # Try to connect to PostgreSQL
+                conn = socket.create_connection(("localhost", self.port), timeout=1)
+                conn.close()
+                # Wait a bit more to ensure PostgreSQL is fully initialized
+                time.sleep(2)
+                logger.info(f"PostgreSQL container is ready after {i + 1} attempt(s)")
+                return True
+            except (socket.error, ConnectionRefusedError) as e:
+                if i == max_retries - 1:
+                    logger.warning(
+                        f"Failed to connect after {max_retries} attempts: {e}"
+                    )
+                    return False
+                time.sleep(retry_interval)
+            except Exception as e:
+                logger.warning(f"Unexpected error checking container readiness: {e}")
+                if i == max_retries - 1:
+                    return False
+                time.sleep(retry_interval)
+
+        return False
+
+    def stop(self) -> None:
+        """
+        Stop the PostgreSQL container.
+
+        This method ensures the container is properly stopped and removed.
+        """
+        if self.container is not None:
+            try:
+                logger.info(f"Stopping PostgreSQL container {self.container_name}")
+                # We've already checked that self.container is not None
+                container = cast(Any, self.container)
+                container.stop(timeout=10)  # Allow 10 seconds for graceful shutdown
+            except Exception as e:
+                logger.warning(f"Failed to stop container: {e}")
+                try:
+                    # Force remove as a fallback
+                    if self.container is not None:
+                        self.container.remove(force=True)
+                    logger.info("Forced container removal")
+                except Exception as e2:
+                    logger.warning(f"Failed to force remove container: {e2}")
+
+    def get_connection_info(self) -> Dict[str, Any]:
+        """
+        Get the connection information for the PostgreSQL container.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing connection parameters
+        """
+        return {
+            "host": "localhost",
+            "port": self.port,
+            "user": self.pg_user,
+            "password": self.pg_password,
+            "database": self.pg_db,
+        }
+
+    def get_connection_string(self) -> str:
+        """
+        Get a PostgreSQL connection string.
+
+        Returns:
+            str: A connection string in the format 'postgresql://user:password@host:port/database'
+        """
+        return f"postgresql://{self.pg_user}:{self.pg_password}@localhost:{self.port}/{self.pg_db}"

{iceaxe-0.7.0.dev3 → iceaxe-0.7.1}/iceaxe/__tests__/migrations/test_action_sorter.py

@@ -19,7 +19,7 @@ class MockNode(DBObject):
     async def create(self, actor: DatabaseActions):
         pass
 
-    async def migrate(self, previous: "MockNode", actor: DatabaseActions):
+    async def migrate(self, previous: DBObject, actor: DatabaseActions):
         pass
 
     async def destroy(self, actor: DatabaseActions):

{iceaxe-0.7.0.dev3 → iceaxe-0.7.1}/iceaxe/__tests__/test_session.py

@@ -780,7 +780,9 @@ async def test_upsert_multiple_conflict_fields(db_connection: DBConnection):
 
 
 @pytest.mark.asyncio
-async def test_for_update_prevents_concurrent_modification(db_connection: DBConnection):
+async def test_for_update_prevents_concurrent_modification(
+    db_connection: DBConnection, docker_postgres
+):
     """
     Test that FOR UPDATE actually locks the row for concurrent modifications.
     """
@@ -799,11 +801,11 @@ async def test_for_update_prevents_concurrent_modification(db_connection: DBConn
         # until our transaction is done
         other_conn = DBConnection(
             await asyncpg.connect(
-                host="localhost",
-                port=5438,
-                user="iceaxe",
-                password="mysecretpassword",
-                database="iceaxe_test_db",
+                host=docker_postgres["host"],
+                port=docker_postgres["port"],
+                user=docker_postgres["user"],
+                password=docker_postgres["password"],
+                database=docker_postgres["database"],
             )
         )
         try:
@@ -820,7 +822,7 @@ async def test_for_update_prevents_concurrent_modification(db_connection: DBConn
 
 
 @pytest.mark.asyncio
-async def test_for_update_skip_locked(db_connection: DBConnection):
+async def test_for_update_skip_locked(db_connection: DBConnection, docker_postgres):
     """
     Test that SKIP LOCKED works as expected.
     """
@@ -844,11 +846,11 @@ async def test_for_update_skip_locked(db_connection: DBConnection):
         # From another connection, try to select both users with SKIP LOCKED
         other_conn = DBConnection(
             await asyncpg.connect(
-                host="localhost",
-                port=5438,
-                user="iceaxe",
-                password="mysecretpassword",
-                database="iceaxe_test_db",
+                host=docker_postgres["host"],
+                port=docker_postgres["port"],
+                user=docker_postgres["user"],
+                password=docker_postgres["password"],
+                database=docker_postgres["database"],
             )
         )
         try:
@@ -866,7 +868,7 @@ async def test_for_update_skip_locked(db_connection: DBConnection):
 
 
 @pytest.mark.asyncio
-async def test_for_update_of_with_join(db_connection: DBConnection):
+async def test_for_update_of_with_join(db_connection: DBConnection, docker_postgres):
     """
     Test FOR UPDATE OF with JOINed tables.
     """
@@ -892,11 +894,11 @@ async def test_for_update_of_with_join(db_connection: DBConnection):
         # but not the artifact
         other_conn = DBConnection(
             await asyncpg.connect(
-                host="localhost",
-                port=5438,
-                user="iceaxe",
-                password="mysecretpassword",
-                database="iceaxe_test_db",
+                host=docker_postgres["host"],
+                port=docker_postgres["port"],
+                user=docker_postgres["user"],
+                password=docker_postgres["password"],
+                database=docker_postgres["database"],
             )
         )
         try:
@@ -1356,7 +1358,7 @@ async def test_batch_upsert_multiple_with_real_db(db_connection: DBConnection):
 
 
 @pytest.mark.asyncio
-async def test_initialize_types_caching():
+async def test_initialize_types_caching(docker_postgres):
     # Clear the global cache for isolation.
     TYPE_CACHE.clear()
 
@@ -1374,11 +1376,11 @@ async def test_initialize_types_caching():
 
     # Establish the first connection.
     conn1 = await asyncpg.connect(
-        host="localhost",
-        port=5438,
-        user="iceaxe",
-        password="mysecretpassword",
-        database="iceaxe_test_db",
+        host=docker_postgres["host"],
+        port=docker_postgres["port"],
+        user=docker_postgres["user"],
+        password=docker_postgres["password"],
+        database=docker_postgres["database"],
     )
     db1 = DBConnection(conn1)
 
@@ -1417,11 +1419,11 @@ async def test_initialize_types_caching():
 
         # Create a second connection to the same database.
         conn2 = await asyncpg.connect(
-            host="localhost",
-            port=5438,
-            user="iceaxe",
-            password="mysecretpassword",
-            database="iceaxe_test_db",
+            host=docker_postgres["host"],
+            port=docker_postgres["port"],
+            user=docker_postgres["user"],
+            password=docker_postgres["password"],
+            database=docker_postgres["database"],
         )
         db2 = DBConnection(conn2)
 
@@ -1475,7 +1477,7 @@ async def test_get_dsn(db_connection: DBConnection):
     assert dsn.startswith("postgresql://")
     assert "iceaxe" in dsn
     assert "localhost" in dsn
-    assert "5438" in dsn
+    assert ":" in dsn  # Just verify there is a port
     assert "iceaxe_test_db" in dsn
 
 

{iceaxe-0.7.0.dev3 → iceaxe-0.7.1}/iceaxe/migrations/migration.py

@@ -1,4 +1,5 @@
 from abc import abstractmethod
+from contextlib import asynccontextmanager
 
 from iceaxe.migrations.migrator import Migrator
 from iceaxe.session import DBConnection
@@ -20,12 +21,23 @@ class MigrationRevisionBase:
     up_revision: str
     down_revision: str | None
 
+    use_transaction: bool = True
+    """
+    Disables the transaction for the current migration. Only do this if you're
+    confident that the migration will succeed on the first try, or is otherwise
+    independent so it can be run multiple times.
+
+    This can speed up migrations, and in some cases might be even fully required for your
+    production database to avoid deadlocks when interacting with hot tables.
+
+    """
+
     async def _handle_up(self, db_connection: DBConnection):
         """
         Internal method to handle the up migration.
         """
         # Isolated migrator context just for this migration
-        async with db_connection.transaction():
+        async with self._optional_transaction(db_connection):
             migrator = Migrator(db_connection)
             await self.up(migrator)
             await migrator.set_active_revision(self.up_revision)
@@ -34,11 +46,19 @@ class MigrationRevisionBase:
         """
         Internal method to handle the down migration.
         """
-        async with db_connection.transaction():
+        async with self._optional_transaction(db_connection):
             migrator = Migrator(db_connection)
             await self.down(migrator)
             await migrator.set_active_revision(self.down_revision)
 
+    @asynccontextmanager
+    async def _optional_transaction(self, db_connection: DBConnection):
+        if self.use_transaction:
+            async with db_connection.transaction():
+                yield
+        else:
+            yield
+
     @abstractmethod
     async def up(self, migrator: Migrator):
         """

{iceaxe-0.7.0.dev3 → iceaxe-0.7.1}/iceaxe/migrations/migrator.py

@@ -86,3 +86,16 @@ class Migrator:
 
         result = await self.db_connection.conn.fetch(query)
         return cast(str | None, result[0]["active_revision"] if result else None)
+
+    async def raw_sql(self, query: str, *args):
+        """
+        Shortcut to execute a raw SQL query against the database. Raw SQL can be more useful
+        than using ORM objects within migrations, because you can interact with the old & new data
+        schemas via text (whereas the runtime ORM is only aware of the current schema).
+
+        ```python {{sticky: True}}
+        await migrator.execute("CREATE TABLE IF NOT EXISTS users (id SERIAL PRIMARY KEY, name VARCHAR(255))")
+        ```
+
+        """
+        await self.db_connection.conn.execute(query, *args)