Support async sessions in strawberry_sqlalchemy_mapper and deprecate sync.

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
catastrophic performance implications. Not only will all resolvers 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.

There's no reason for us to support a foot bazooka, hence the deprecation.

Author: mattalbr

RELEASE.md

@@ -0,0 +1,3 @@
+Release type: minor
+
+Adds support for async sessions and deprecates sync sessions due to performance reasons.

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,10 @@
+import logging
 from collections import defaultdict
-from typing import Any, Dict, List, Mapping, Tuple, Union
+from typing import Any, 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 +16,37 @@ 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[
+            Callable[[], Union[AsyncSession, 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."
+            )
+        if bind is not None:
+            # For anyone coming here because of this warning:
+            # Making blocking database calls from within an async function (the resolver) has
+            # catastrophic performance implications. Not only will all resolvers 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.
+            self._logger.warning(
+                "`bind` parameter is deprecated due to performance issues. Use `async_bind_factory` instead."
+            )
+
+    async def _scalars(self, *args, **kwargs):
+        if self._async_bind_factory:
+            return await self._async_bind_factory().scalars(*args, **kwargs)
+        else:
+            # Deprecated, but supported for now.
+            return self._bind.scalars(*args, **kwargs)
 
     def loader_for(self, relationship: RelationshipProperty) -> DataLoader:
         """
@@ -35,7 +65,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