Support async sessions in strawberry_sqlalchemy_mapper (#53)

I left a note in the loader's init, but to duplicate that here: Making blocking database calls from within an async function (the resolver) has potentially catastrophic performance implications. Not only will all field resolution be effectively serialized, any other coroutines waiting on the event loop (e.g. concurrent requests in a web server), will be blocked as well, grinding your entire service to a halt.

More discussion needs to happen about what we do with the sync bind.

Author: mattalbr

RELEASE.md

@@ -0,0 +1,14 @@
+Release type: minor
+
+Adds support for async sessions. To use:
+
+```python
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+from strawberry_sqlalchemy_mapper import StrawberrySQLAlchemyLoader
+
+url = "postgresql://..."
+engine = create_async_engine(url)
+sessionmaker = async_sessionmaker(engine)
+
+loader = StrawberrySQLAlchemyLoader(async_bind_factory=sessionmaker)
+```

poetry.lock

@@ -40,12 +40,13 @@ version_scheme = "no-guess-dev"
 
 [tool.poetry.dependencies]
 python = "^3.8"
-sqlalchemy = ">=1.4"
+sqlalchemy = {extras = ["asyncio"], version = ">=1.4"}
 strawberry-graphql = ">=0.95"
 sentinel = ">=0.3,<1.1"
 greenlet = {version = ">=3.0.0rc1", python = ">=3.12"}
 
 [tool.poetry.group.dev.dependencies]
+asyncpg = "^0.28.0"
 black = ">=22,<24"
 importlib-metadata = ">=4.11.1,<7.0.0"
 mypy = "1.5.1"
@@ -106,6 +107,8 @@ ignore = [
     # we'd want to have consistent docstrings in future
     "D",
     "ANN001", # missing annotation for function argument self.
+    "ANN002", # missing annotation for *args.
+    "ANN003", # missing annotatino for **kwargs.
     "ANN101", # missing annotation for self?
     # definitely enable these, maybe not in tests
     "ANN102",

pyproject.toml

@@ -1,8 +1,20 @@
+import logging
 from collections import defaultdict
-from typing import Any, Dict, List, Mapping, Tuple, Union
+from typing import (
+    Any,
+    AsyncContextManager,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+)
 
 from sqlalchemy import select, tuple_
 from sqlalchemy.engine.base import Connection
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 from sqlalchemy.orm import RelationshipProperty, Session
 from strawberry.dataloader import DataLoader
 
@@ -14,9 +26,32 @@ class StrawberrySQLAlchemyLoader:
 
     _loaders: Dict[RelationshipProperty, DataLoader]
 
-    def __init__(self, bind: Union[Session, Connection]) -> None:
+    def __init__(
+        self,
+        bind: Union[Session, Connection, None] = None,
+        async_bind_factory: Optional[
+            Union[
+                Callable[[], AsyncContextManager[AsyncSession]],
+                Callable[[], AsyncContextManager[AsyncConnection]],
+            ]
+        ] = None,
+    ) -> None:
         self._loaders = {}
-        self.bind = bind
+        self._bind = bind
+        self._async_bind_factory = async_bind_factory
+        self._logger = logging.getLogger("strawberry_sqlalchemy_mapper")
+        if bind is None and async_bind_factory is None:
+            self._logger.warning(
+                "One of bind or async_bind_factory must be set for loader to function properly."
+            )
+
+    async def _scalars(self, *args, **kwargs):
+        if self._async_bind_factory:
+            async with self._async_bind_factory() as bind:
+                return await bind.scalars(*args, **kwargs)
+        else:
+            assert self._bind is not None
+            return self._bind.scalars(*args, **kwargs)
 
     def loader_for(self, relationship: RelationshipProperty) -> DataLoader:
         """
@@ -35,7 +70,7 @@ async def load_fn(keys: List[Tuple]) -> List[Any]:
                 )
                 if relationship.order_by:
                     query = query.order_by(*relationship.order_by)
-                rows = self.bind.scalars(query).all()
+                rows = (await self._scalars(query)).all()
 
                 def group_by_remote_key(row: Any) -> Tuple:
                     return tuple(

src/strawberry_sqlalchemy_mapper/loader.py

@@ -17,6 +17,9 @@
 from packaging import version
 from sqlalchemy import orm
 from sqlalchemy.engine import Engine
+from sqlalchemy.ext import asyncio
+from sqlalchemy.ext.asyncio import AsyncAttrs, create_async_engine
+from sqlalchemy.ext.asyncio.engine import AsyncEngine
 from testing.postgresql import Postgresql, PostgresqlFactory
 
 SQLA_VERSION = version.parse(sqlalchemy.__version__)
@@ -57,7 +60,11 @@ def postgresql(postgresql_factory) -> Postgresql:
 @pytest.fixture(params=SUPPORTED_DBS)
 def engine(request) -> Engine:
     if request.param == "postgresql":
-        url = request.getfixturevalue("postgresql").url()
+        url = (
+            request.getfixturevalue("postgresql")
+            .url()
+            .replace("postgresql://", "postgresql+psycopg2://")
+        )
     else:
         raise ValueError("Unsupported database: %s", request.param)
     kwargs = {}
@@ -72,6 +79,31 @@ def sessionmaker(engine) -> orm.sessionmaker:
     return orm.sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 
+@pytest.fixture(params=SUPPORTED_DBS)
+def async_engine(request) -> AsyncEngine:
+    if request.param == "postgresql":
+        url = (
+            request.getfixturevalue("postgresql")
+            .url()
+            .replace("postgresql://", "postgresql+asyncpg://")
+        )
+    else:
+        raise ValueError("Unsupported database: %s", request.param)
+    kwargs = {}
+    if not SQLA2:
+        kwargs["future"] = True
+    engine = create_async_engine(url, **kwargs)
+    return engine
+
+
+@pytest.fixture
+def async_sessionmaker(async_engine) -> asyncio.async_sessionmaker:
+    return asyncio.async_sessionmaker(async_engine)
+
+
 @pytest.fixture
 def base():
-    return orm.declarative_base()
+    class Base(AsyncAttrs, orm.DeclarativeBase):
+        pass
+
+    return Base

tests/conftest.py

@@ -58,7 +58,7 @@ class Department(base):
 
 def test_loader_init():
     loader = StrawberrySQLAlchemyLoader(bind=None)
-    assert loader.bind is None
+    assert loader._bind is None
     assert loader._loaders == {}
 
 
@@ -99,14 +99,49 @@ async def test_loader_for(engine, base, sessionmaker, many_to_one_tables):
         assert department.name == "d2"
 
         loader = base_loader.loader_for(Department.employees.property)
-        key = tuple(
+
+        employees = await loader.load((d2.id,))
+        assert {e.name for e in employees} == {"e1"}
+
+
+@pytest.mark.asyncio
+async def test_loader_with_async_session(
+    async_engine, base, async_sessionmaker, many_to_one_tables
+):
+    Employee, Department = many_to_one_tables
+    async with async_engine.begin() as conn:
+        await conn.run_sync(base.metadata.create_all)
+
+    async with async_sessionmaker() as session:
+        e1 = Employee(name="e1")
+        e2 = Employee(name="e2")
+        d1 = Department(name="d1")
+        d2 = Department(name="d2")
+        session.add(e1)
+        session.add(e2)
+        session.add(d1)
+        session.add(d2)
+        await session.flush()
+
+        e1.department = d2
+        e2.department = d1
+        await session.commit()
+        d2_id = await d2.awaitable_attrs.id
+        department_loader_key = tuple(
             [
-                getattr(d2, local.key)
-                for local, _ in Department.employees.property.local_remote_pairs
+                await getattr(e1.awaitable_attrs, local.key)
+                for local, _ in Employee.department.property.local_remote_pairs
             ]
         )
-        employees = await loader.load((d2.id,))
-        assert {e.name for e in employees} == {"e1"}
+    base_loader = StrawberrySQLAlchemyLoader(async_bind_factory=async_sessionmaker)
+    loader = base_loader.loader_for(Employee.department.property)
+
+    department = await loader.load(department_loader_key)
+    assert department.name == "d2"
+
+    loader = base_loader.loader_for(Department.employees.property)
+    employees = await loader.load((d2_id,))
+    assert {e.name for e in employees} == {"e1"}
 
 
 @pytest.mark.xfail