feat(async): add support for async sessions

graphene_sqlalchemy/fields.py

@@ -11,7 +11,10 @@
 from graphql_relay import connection_from_array_slice
 
 from .batching import get_batch_resolver
-from .utils import EnumValue, get_query
+from .utils import EnumValue, get_query, get_session, is_sqlalchemy_version_less_than
+
+if not is_sqlalchemy_version_less_than("1.4"):
+    from sqlalchemy.ext.asyncio import AsyncSession
 
 
 class SQLAlchemyConnectionField(ConnectionField):
@@ -81,8 +84,51 @@ def get_query(cls, model, info, sort=None, **args):
 
     @classmethod
     def resolve_connection(cls, connection_type, model, info, args, resolved):
+        session = get_session(info.context)
+        if resolved is None:
+            if not is_sqlalchemy_version_less_than("1.4") and isinstance(
+                session, AsyncSession
+            ):
+
+                async def get_result():
+                    return await cls.resolve_connection_async(
+                        connection_type, model, info, args, resolved
+                    )
+
+                return get_result()
+
+            else:
+                resolved = cls.get_query(model, info, **args)
+        if isinstance(resolved, Query):
+            _len = resolved.count()
+        else:
+            _len = len(resolved)
+
+        def adjusted_connection_adapter(edges, pageInfo):
+            return connection_adapter(connection_type, edges, pageInfo)
+
+        connection = connection_from_array_slice(
+            array_slice=resolved,
+            args=args,
+            slice_start=0,
+            array_length=_len,
+            array_slice_length=_len,
+            connection_type=adjusted_connection_adapter,
+            edge_type=connection_type.Edge,
+            page_info_type=page_info_adapter,
+        )
+        connection.iterable = resolved
+        connection.length = _len
+        return connection
+
+    @classmethod
+    async def resolve_connection_async(
+        cls, connection_type, model, info, args, resolved
+    ):
+        session = get_session(info.context)
         if resolved is None:
-            resolved = cls.get_query(model, info, **args)
+            query = cls.get_query(model, info, **args)
+            resolved = (await session.scalars(query)).all()
         if isinstance(resolved, Query):
             _len = resolved.count()
         else:
@@ -179,7 +225,7 @@ def from_relationship(cls, relationship, registry, **field_kwargs):
         return cls(
             model_type.connection,
             resolver=get_batch_resolver(relationship),
-            **field_kwargs
+            **field_kwargs,
         )
 
 

graphene_sqlalchemy/tests/conftest.py

@@ -1,14 +1,17 @@
 import pytest
+import pytest_asyncio
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
 import graphene
+from graphene_sqlalchemy.utils import is_sqlalchemy_version_less_than
 
 from ..converter import convert_sqlalchemy_composite
 from ..registry import reset_global_registry
 from .models import Base, CompositeFullName
 
-test_db_url = "sqlite://"  # use in-memory database for tests
+if not is_sqlalchemy_version_less_than("1.4"):
+    from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 
 
 @pytest.fixture(autouse=True)
@@ -22,18 +25,49 @@ def convert_composite_class(composite, registry):
         return graphene.Field(graphene.Int)
 
 
-@pytest.fixture(scope="function")
-def session_factory():
-    engine = create_engine(test_db_url)
-    Base.metadata.create_all(engine)
+@pytest.fixture(params=[False, True])
+def async_session(request):
+    return request.param
+
+
+@pytest.fixture
+def test_db_url(async_session: bool):
+    if async_session:
+        return "sqlite+aiosqlite://"
+    else:
+        return "sqlite://"
 
-    yield sessionmaker(bind=engine)
 
+@pytest.mark.asyncio
+@pytest_asyncio.fixture(scope="function")
+async def session_factory(async_session: bool, test_db_url: str):
+    if async_session:
+        if is_sqlalchemy_version_less_than("1.4"):
+            pytest.skip("Async Sessions only work in sql alchemy 1.4 and above")
+        engine = create_async_engine(test_db_url)
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
+        yield sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False)
+        await engine.dispose()
+    else:
+        engine = create_engine(test_db_url)
+        Base.metadata.create_all(engine)
+        yield sessionmaker(bind=engine, expire_on_commit=False)
+        # SQLite in-memory db is deleted when its connection is closed.
+        # https://www.sqlite.org/inmemorydb.html
+        engine.dispose()
+
+
+@pytest_asyncio.fixture(scope="function")
+async def sync_session_factory():
+    engine = create_engine("sqlite://")
+    Base.metadata.create_all(engine)
+    yield sessionmaker(bind=engine, expire_on_commit=False)
     # SQLite in-memory db is deleted when its connection is closed.
     # https://www.sqlite.org/inmemorydb.html
     engine.dispose()
 
 
-@pytest.fixture(scope="function")
+@pytest_asyncio.fixture(scope="function")
 def session(session_factory):
     return session_factory()

graphene_sqlalchemy/tests/models.py

@@ -75,10 +75,14 @@ class Reporter(Base):
     email = Column(String(), doc="Email")
     favorite_pet_kind = Column(PetKind)
     pets = relationship(
-        "Pet", secondary=association_table, backref="reporters", order_by="Pet.id"
+        "Pet",
+        secondary=association_table,
+        backref="reporters",
+        order_by="Pet.id",
+        lazy="selectin",
     )
-    articles = relationship("Article", backref="reporter")
-    favorite_article = relationship("Article", uselist=False)
+    articles = relationship("Article", backref="reporter", lazy="selectin")
+    favorite_article = relationship("Article", uselist=False, lazy="selectin")
 
     @hybrid_property
     def hybrid_prop_with_doc(self):

graphene_sqlalchemy/tests/models_batching.py

@@ -0,0 +1,91 @@
+from __future__ import absolute_import
+
+import enum
+
+from sqlalchemy import (
+    Column,
+    Date,
+    Enum,
+    ForeignKey,
+    Integer,
+    String,
+    Table,
+    func,
+    select,
+)
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import column_property, relationship
+
+PetKind = Enum("cat", "dog", name="pet_kind")
+
+
+class HairKind(enum.Enum):
+    LONG = "long"
+    SHORT = "short"
+
+
+Base = declarative_base()
+
+association_table = Table(
+    "association",
+    Base.metadata,
+    Column("pet_id", Integer, ForeignKey("pets.id")),
+    Column("reporter_id", Integer, ForeignKey("reporters.id")),
+)
+
+
+class Pet(Base):
+    __tablename__ = "pets"
+    id = Column(Integer(), primary_key=True)
+    name = Column(String(30))
+    pet_kind = Column(PetKind, nullable=False)
+    hair_kind = Column(Enum(HairKind, name="hair_kind"), nullable=False)
+    reporter_id = Column(Integer(), ForeignKey("reporters.id"))
+
+
+class Reporter(Base):
+    __tablename__ = "reporters"
+
+    id = Column(Integer(), primary_key=True)
+    first_name = Column(String(30), doc="First name")
+    last_name = Column(String(30), doc="Last name")
+    email = Column(String(), doc="Email")
+    favorite_pet_kind = Column(PetKind)
+    pets = relationship(
+        "Pet",
+        secondary=association_table,
+        backref="reporters",
+        order_by="Pet.id",
+    )
+    articles = relationship("Article", backref="reporter")
+    favorite_article = relationship("Article", uselist=False)
+
+    column_prop = column_property(
+        select([func.cast(func.count(id), Integer)]), doc="Column property"
+    )
+
+
+class Article(Base):
+    __tablename__ = "articles"
+    id = Column(Integer(), primary_key=True)
+    headline = Column(String(100))
+    pub_date = Column(Date())
+    reporter_id = Column(Integer(), ForeignKey("reporters.id"))
+    readers = relationship(
+        "Reader", secondary="articles_readers", back_populates="articles"
+    )
+
+
+class Reader(Base):
+    __tablename__ = "readers"
+    id = Column(Integer(), primary_key=True)
+    name = Column(String(100))
+    articles = relationship(
+        "Article", secondary="articles_readers", back_populates="readers"
+    )
+
+
+class ArticleReader(Base):
+    __tablename__ = "articles_readers"
+    article_id = Column(Integer(), ForeignKey("articles.id"), primary_key=True)
+    reader_id = Column(Integer(), ForeignKey("readers.id"), primary_key=True)