Merge pull request #32 from soapun/feature/async-sqla

Async SQLAlchemy Support

Author: wey-gu

src/py_pglite/pytest_plugin.py

@@ -89,9 +89,14 @@ def _should_disable_django_plugin(config: pytest.Config) -> bool:
 # Smart fixture loading with perfect isolation
 if HAS_SQLALCHEMY:
     try:
+        # async
+        from py_pglite.sqlalchemy.fixtures import pglite_async_engine
+        from py_pglite.sqlalchemy.fixtures import pglite_async_session
+        from py_pglite.sqlalchemy.fixtures import pglite_async_sqlalchemy_manager
         from py_pglite.sqlalchemy.fixtures import pglite_config
         from py_pglite.sqlalchemy.fixtures import pglite_engine
         from py_pglite.sqlalchemy.fixtures import pglite_session
+        from py_pglite.sqlalchemy.fixtures import pglite_sqlalchemy_async_engine
         from py_pglite.sqlalchemy.fixtures import pglite_sqlalchemy_engine
         from py_pglite.sqlalchemy.fixtures import pglite_sqlalchemy_manager
         from py_pglite.sqlalchemy.fixtures import pglite_sqlalchemy_session

src/py_pglite/sqlalchemy/__init__.py

@@ -4,26 +4,36 @@
 """
 
 # Import fixtures and utilities
+from py_pglite.sqlalchemy.fixtures import pglite_async_engine
+from py_pglite.sqlalchemy.fixtures import pglite_async_session
+from py_pglite.sqlalchemy.fixtures import pglite_async_sqlalchemy_manager
 from py_pglite.sqlalchemy.fixtures import pglite_engine
 from py_pglite.sqlalchemy.fixtures import pglite_session
+from py_pglite.sqlalchemy.fixtures import pglite_sqlalchemy_async_engine
 from py_pglite.sqlalchemy.fixtures import pglite_sqlalchemy_engine
 from py_pglite.sqlalchemy.fixtures import pglite_sqlalchemy_session
 from py_pglite.sqlalchemy.manager import SQLAlchemyPGliteManager
+from py_pglite.sqlalchemy.manager_async import SQLAlchemyAsyncPGliteManager
 from py_pglite.sqlalchemy.utils import create_all_tables
 from py_pglite.sqlalchemy.utils import drop_all_tables
 from py_pglite.sqlalchemy.utils import get_session_class
 
 
 __all__ = [
     # Manager
+    "SQLAlchemyAsyncPGliteManager",
     "SQLAlchemyPGliteManager",
     # Utilities
     "create_all_tables",
     "drop_all_tables",
     "get_session_class",
     # Fixtures
+    "pglite_async_engine",
+    "pglite_async_session",
+    "pglite_async_sqlalchemy_manager",
     "pglite_engine",
     "pglite_session",
+    "pglite_sqlalchemy_async_engine",
     "pglite_sqlalchemy_engine",
     "pglite_sqlalchemy_session",
 ]

src/py_pglite/sqlalchemy/fixtures.py

@@ -3,24 +3,30 @@
 import logging
 import time
 
+from collections.abc import AsyncGenerator
 from collections.abc import Generator
 from typing import Any
 
 import pytest
 
 from sqlalchemy import text
 from sqlalchemy.engine import Engine
+from sqlalchemy.ext.asyncio import AsyncEngine
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.ext.asyncio import async_sessionmaker
 from sqlalchemy.orm import Session
 from sqlalchemy.orm import sessionmaker
 
 from py_pglite.config import PGliteConfig
 from py_pglite.sqlalchemy.manager import SQLAlchemyPGliteManager
+from py_pglite.sqlalchemy.manager_async import SQLAlchemyAsyncPGliteManager
 
 
 # Try to import SQLModel
 try:
     from sqlmodel import Session as SQLModelSession
     from sqlmodel import SQLModel
+    from sqlmodel.ext.asyncio.session import AsyncSession as SQLModelAsyncSession
 
     HAS_SQLMODEL = True
 except ImportError:
@@ -55,6 +61,25 @@ def pglite_sqlalchemy_manager(
         manager.stop()
 
 
+@pytest.fixture(scope="session")
+async def pglite_async_sqlalchemy_manager(
+    pglite_config: PGliteConfig,
+) -> AsyncGenerator[SQLAlchemyAsyncPGliteManager, None, None]:
+    """Pytest fixture providing an async SQLAlchemy-enabled PGlite manager."""
+    manager = SQLAlchemyAsyncPGliteManager(pglite_config)
+    manager.start()
+
+    is_ready = await manager.wait_for_ready()
+    # Wait for database to be ready
+    if not is_ready:
+        raise RuntimeError("Failed to start PGlite database")
+
+    try:
+        yield manager
+    finally:
+        await manager.stop()
+
+
 @pytest.fixture(scope="session")
 def pglite_engine(pglite_sqlalchemy_manager: SQLAlchemyPGliteManager) -> Engine:
     """Pytest fixture providing a SQLAlchemy engine connected to PGlite.
@@ -65,6 +90,18 @@ def pglite_engine(pglite_sqlalchemy_manager: SQLAlchemyPGliteManager) -> Engine:
     return pglite_sqlalchemy_manager.get_engine()
 
 
+@pytest.fixture(scope="session")
+def pglite_async_engine(
+    pglite_async_sqlalchemy_manager: SQLAlchemyAsyncPGliteManager,
+) -> AsyncEngine:
+    """Pytest fixture providing a SQLAlchemy async engine connected to PGlite.
+
+    Uses the SQLAlchemy-enabled manager to ensure proper SQLAlchemy integration.
+    """
+    # Use the shared engine from manager (no custom parameters to avoid conflicts)
+    return pglite_async_sqlalchemy_manager.get_engine()
+
+
 @pytest.fixture(scope="session")
 def pglite_sqlalchemy_engine(
     pglite_sqlalchemy_manager: SQLAlchemyPGliteManager,
@@ -74,6 +111,15 @@ def pglite_sqlalchemy_engine(
     return pglite_sqlalchemy_manager.get_engine()
 
 
+@pytest.fixture(scope="session")
+def pglite_sqlalchemy_async_engine(
+    pglite_async_sqlalchemy_manager: SQLAlchemyAsyncPGliteManager,
+) -> AsyncEngine:
+    """Pytest fixture providing an optimized SQLAlchemy async engine connected to PGlite."""
+    # Use the shared engine from manager (no custom parameters to avoid conflicts)
+    return pglite_async_sqlalchemy_manager.get_engine()
+
+
 @pytest.fixture(scope="function")
 def pglite_session(pglite_engine: Engine) -> Generator[Any, None, None]:
     """Pytest fixture providing a SQLAlchemy/SQLModel session with proper isolation.
@@ -161,6 +207,96 @@ def pglite_session(pglite_engine: Engine) -> Generator[Any, None, None]:
             logger.warning(f"Error closing session: {e}")
 
 
+@pytest.fixture(scope="function")
+async def pglite_async_session(
+    pglite_async_engine: AsyncEngine,
+) -> AsyncGenerator[Any, None, None]:
+    """Pytest fixture providing a SQLAlchemy/SQLModel async session with proper isolation.
+
+    This fixture ensures database isolation between tests by cleaning all data
+    at the start of each test.
+    """
+    # Clean up data before test starts
+    logger.info("Starting database cleanup before test...")
+    retry_count = 3
+    for attempt in range(retry_count):
+        try:
+            async with pglite_async_engine.connect() as conn:
+                # Get all table names from information_schema
+                result = await conn.execute(
+                    text("""
+                    SELECT table_name
+                    FROM information_schema.tables
+                    WHERE table_schema = 'public'
+                    AND table_type = 'BASE TABLE'
+                """)
+                )
+
+                table_names = [row[0] for row in result]
+                logger.info(f"Found tables to clean: {table_names}")
+
+                if table_names:
+                    # Disable foreign key checks for faster cleanup
+                    await conn.execute(text("SET session_replication_role = replica;"))
+
+                    # Truncate all tables
+                    for table_name in table_names:
+                        logger.info(f"Truncating table: {table_name}")
+                        await conn.execute(
+                            text(
+                                f'TRUNCATE TABLE "{table_name}" '
+                                f"RESTART IDENTITY CASCADE;"
+                            )
+                        )
+
+                    # Re-enable foreign key checks
+                    await conn.execute(text("SET session_replication_role = DEFAULT;"))
+
+                    # Commit the cleanup
+                    await conn.commit()
+                    logger.info("Database cleanup completed successfully")
+                else:
+                    logger.info("No tables found to clean")
+                break  # Success, exit retry loop
+
+        except Exception as e:
+            logger.info(f"Database cleanup attempt {attempt + 1} failed: {e}")
+            if attempt == retry_count - 1:
+                logger.warning(
+                    "Database cleanup failed after all retries, continuing anyway"
+                )
+            else:
+                time.sleep(0.5)  # Brief pause before retry
+
+    # Create session - prefer SQLModel if available
+    if HAS_SQLMODEL and SQLModelAsyncSession is not None:
+        session = SQLModelAsyncSession(pglite_async_engine)
+        # Create tables if using SQLModel with retry logic
+        if SQLModel is not None:
+            for attempt in range(3):
+                try:
+                    async with pglite_async_engine.begin() as conn:
+                        await conn.run_sync(SQLModel.metadata.create_all)
+                    break
+                except Exception as e:
+                    logger.warning(f"Table creation attempt {attempt + 1} failed: {e}")
+                    if attempt == 2:
+                        raise
+                    time.sleep(0.5)
+    else:
+        session_local = async_sessionmaker(bind=pglite_async_engine)
+        session = session_local()  # type: ignore[assignment]
+
+    try:
+        yield session
+    finally:
+        # Close the session safely
+        try:
+            await session.close()
+        except Exception as e:
+            logger.warning(f"Error closing session: {e}")
+
+
 @pytest.fixture(scope="function")
 def pglite_sqlalchemy_session(pglite_session: Session) -> Session:
     """Legacy fixture name for backwards compatibility."""