feat: sqlalchemy_url accepts async dialect - DIA-46202 (#62)

Author: hadrien

README.md

@@ -92,7 +92,7 @@ EOF
 
 Running the app:
 ```bash
-sqlalchemy_url=sqlite:///db.sqlite?check_same_thread=false uvicorn main:app --reload
+sqlalchemy_url=sqlite:///db.sqlite?check_same_thread=false uvicorn main:app
 ```
 
 # Configuration
@@ -119,10 +119,10 @@ To enable `asyncio` support against a Postgres DB, install `asyncpg`:
 pip install asyncpg
 ```
 
-And define environment variable `async_sqlalchemy_url` with `postgres+asyncpg` scheme:
+And define environment variable `sqlalchemy_url` with `postgres+asyncpg` scheme:
 
 ```bash
-export async_sqlalchemy_url=postgresql+asyncpg://postgres@localhost
+export sqlalchemy_url=postgresql+asyncpg://postgres@localhost
 ```
 
 ## Setup the app:

fastapi_sqla/__init__.py

@@ -12,6 +12,7 @@
 from pydantic import BaseModel, Field
 from pydantic.generics import GenericModel
 from sqlalchemy import engine_from_config, text
+from sqlalchemy.engine import Engine
 from sqlalchemy.ext.declarative import DeferredReflection
 from sqlalchemy.orm import Query as LegacyQuery
 from sqlalchemy.orm.session import Session as SqlaSession
@@ -51,38 +52,48 @@
 _Session = sessionmaker()
 
 
+def new_engine(*, envvar_prefix: str = None) -> Engine:
+    envvar_prefix = envvar_prefix if envvar_prefix else "sqlalchemy_"
+    lowercase_environ = {
+        k.lower(): v for k, v in os.environ.items() if k.lower() != "sqlalchemy_warn_20"
+    }
+    return engine_from_config(lowercase_environ, prefix=envvar_prefix)
+
+
+def is_async_dialect(engine):
+    return engine.dialect.is_async if hasattr(engine.dialect, "is_async") else False
+
+
 def setup(app: FastAPI):
-    app.add_event_handler("startup", startup)
-    app.middleware("http")(add_session_to_request)
+    engine = new_engine()
+
+    if not is_async_dialect(engine):
+        app.add_event_handler("startup", startup)
+        app.middleware("http")(add_session_to_request)
 
-    async_sqlalchemy_url = os.getenv("async_sqlalchemy_url")
-    if async_sqlalchemy_url:
+    has_async_config = "async_sqlalchemy_url" in os.environ or is_async_dialect(engine)
+    if has_async_config:
         assert asyncio_support, asyncio_support_err
         app.add_event_handler("startup", asyncio_support.startup)
         app.middleware("http")(asyncio_support.add_session_to_request)
 
 
 def startup():
-    lowercase_environ = {
-        k.lower(): v for k, v in os.environ.items() if k.lower() != "sqlalchemy_warn_20"
-    }
-    engine = engine_from_config(lowercase_environ, prefix="sqlalchemy_")
+    engine = new_engine()
     aws_rds_iam_support.setup(engine.engine)
 
-    Base.metadata.bind = engine
-    Base.prepare(engine)
-    _Session.configure(bind=engine)
-
     # Fail early:
     try:
-        with open_session() as session:
-            session.execute(text("select 'OK'"))
+        with engine.connect() as connection:
+            connection.execute(text("select 'OK'"))
     except Exception:
         logger.critical(
             "Fail querying db: is sqlalchemy_url envvar correctly configured?"
         )
         raise
 
+    Base.prepare(engine)
+    _Session.configure(bind=engine)
     logger.info("startup", engine=engine)
 
 

fastapi_sqla/_pytest_plugin.py

@@ -162,29 +162,23 @@ async def async_sqla_connection(async_engine, event_loop):
             await connection.rollback()
 
     @fixture
-    async def patch_create_async_engine(
-        async_sqlalchemy_url, async_sqla_connection, request
-    ):
+    async def patch_new_engine(async_sqlalchemy_url, async_sqla_connection, request):
         """So that all async DB operations are never written to db for real."""
         from fastapi_sqla.asyncio_support import _AsyncSession
 
         if "dont_patch_engines" in request.keywords:
             yield
 
         else:
-            with patch(
-                "fastapi_sqla.asyncio_support.create_async_engine"
-            ) as create_async_engine:
-                create_async_engine.return_value = async_sqla_connection
+            with patch("fastapi_sqla.asyncio_support.new_engine") as new_engine:
+                new_engine.return_value = async_sqla_connection
                 _AsyncSession.configure(
                     bind=async_sqla_connection, expire_on_commit=False
                 )
-                yield create_async_engine
+                yield new_engine
 
     @fixture
-    async def async_session(
-        async_sqla_connection, sqla_reflection, patch_create_async_engine
-    ):
+    async def async_session(async_sqla_connection, sqla_reflection, patch_new_engine):
         from fastapi_sqla.asyncio_support import _AsyncSession
 
         session = _AsyncSession(bind=async_sqla_connection)
@@ -194,5 +188,5 @@ async def async_session(
 else:
 
     @fixture
-    async def patch_create_async_engine():
+    async def patch_new_engine():
         pass

fastapi_sqla/asyncio_support.py

@@ -4,33 +4,45 @@
 import structlog
 from fastapi import Request
 from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import AsyncSession as SqlaAsyncSession
-from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.orm.session import sessionmaker
 
-from . import aws_rds_iam_support
+from . import Base, aws_rds_iam_support, new_engine
 
 logger = structlog.get_logger(__name__)
 _ASYNC_SESSION_KEY = "fastapi_sqla_async_session"
 _AsyncSession = sessionmaker(class_=SqlaAsyncSession)
 
 
+def new_async_engine():
+    envvar_prefix = None
+    if "async_sqlalchemy_url" in os.environ:
+        envvar_prefix = "async_sqlalchemy_"
+
+    engine = new_engine(envvar_prefix=envvar_prefix)
+    return AsyncEngine(engine)
+
+
 async def startup():
-    async_sqlalchemy_url = os.environ["async_sqlalchemy_url"]
-    engine = create_async_engine(async_sqlalchemy_url)
+    engine = new_async_engine()
     aws_rds_iam_support.setup(engine.sync_engine)
-    _AsyncSession.configure(bind=engine, expire_on_commit=False)
 
     # Fail early:
     try:
-        async with open_session() as session:
-            await session.execute(text("select 'ok'"))
+        async with engine.connect() as connection:
+            await connection.execute(text("select 'ok'"))
     except Exception:
         logger.critical(
-            "Fail querying db: is async_sqlalchemy_url envvar correctly configured?"
+            "Failed querying db: is sqlalchemy_url or async_sqlalchemy_url envvar "
+            "correctly configured?"
         )
         raise
 
+    async with engine.connect() as connection:
+        await connection.run_sync(Base.prepare)
+
+    _AsyncSession.configure(bind=engine, expire_on_commit=False)
     logger.info("startup", async_engine=engine)
 
 

tests/test_asyncio_support.py

@@ -1,10 +1,12 @@
-from pytest import fixture, mark
+from unittest.mock import AsyncMock, patch
+
+from pytest import fixture, mark, raises
 from sqlalchemy import text
 
 pytestmark = [mark.sqlalchemy("1.4"), mark.require_asyncpg]
 
 
-@fixture(autouse=True)
+@fixture
 async def setup(environ):
     from fastapi_sqla.asyncio_support import startup
 
@@ -23,10 +25,40 @@ async def test_startup_configure_async_session():
     assert res.scalar() == 123
 
 
-async def test_open_async_session():
+async def test_open_async_session(setup):
     from fastapi_sqla.asyncio_support import open_session
 
     async with open_session() as session:
         res = await session.execute(text("select 123"))
 
     assert res.scalar() == 123
+
+
+async def test_new_async_engine_without_async_alchemy_url(
+    monkeypatch, async_sqlalchemy_url
+):
+    from fastapi_sqla.asyncio_support import new_async_engine
+
+    monkeypatch.delenv("async_sqlalchemy_url")
+    monkeypatch.setenv("sqlalchemy_url", async_sqlalchemy_url)
+
+    assert new_async_engine()
+
+
+@fixture
+def AsyncSessionMock():
+    with patch("fastapi_sqla.asyncio_support._AsyncSession") as AsyncSessionMock:
+        AsyncSessionMock.return_value = AsyncMock()
+        yield AsyncSessionMock
+
+
+async def test_context_manager_rollbacks_on_error(AsyncSessionMock):
+    from fastapi_sqla.asyncio_support import open_session
+
+    session = AsyncSessionMock.return_value
+    with raises(Exception) as raise_info:
+        async with open_session():
+            raise Exception("boom!")
+
+    session.rollback.assert_awaited_once_with()
+    assert raise_info.value.args == ("boom!",)

tests/test_setup.py

@@ -0,0 +1,23 @@
+from unittest.mock import Mock
+
+from pytest import mark
+
+
+@mark.sqlalchemy("1.4")
+@mark.require_asyncpg
+def test_setup_with_async_sqlalchemy_url_adds_asyncio_support_startup(
+    monkeypatch, async_sqlalchemy_url
+):
+    from fastapi_sqla import asyncio_support, setup
+
+    monkeypatch.delenv("async_sqlalchemy_url")
+    monkeypatch.setenv("sqlalchemy_url", async_sqlalchemy_url)
+
+    app = Mock()
+    setup(app)
+
+    app.add_event_handler.assert_called_once_with("startup", asyncio_support.startup)
+    app.middleware.assert_called_once_with("http")
+    app.middleware.return_value.assert_called_once_with(
+        asyncio_support.add_session_to_request
+    )

tests/test_startup.py

@@ -6,7 +6,7 @@
 from pytest import fixture, mark, raises
 from sqlalchemy import text
 
-pytestmark = mark.usefixtures("patch_engine_from_config", "patch_create_async_engine")
+pytestmark = mark.usefixtures("patch_engine_from_config", "patch_new_engine")
 
 
 @fixture(params=[True, False])
@@ -50,6 +50,7 @@ def test_startup(case_sensitive_environ):
     assert session.execute(text("SELECT 1")).scalar() == 1
 
 
+@mark.dont_patch_engines
 async def test_fastapi_integration():
     from fastapi_sqla import _Session, setup