feat: SQLAlchemy 2.0 support (#368)

This PR updates the dataloader and unit tests to be compatible with sqlalchemy 2.0

Author: erikwrede

.github/workflows/tests.yml

@@ -14,8 +14,8 @@ jobs:
     strategy:
       max-parallel: 10
       matrix:
-        sql-alchemy: ["1.2", "1.3", "1.4"]
-        python-version: ["3.7", "3.8", "3.9", "3.10"]
+        sql-alchemy: [ "1.2", "1.3", "1.4","2.0" ]
+        python-version: [ "3.7", "3.8", "3.9", "3.10" ]
 
     steps:
       - uses: actions/checkout@v3

.gitignore

@@ -12,6 +12,7 @@ __pycache__/
 .Python
 env/
 .venv/
+venv/
 build/
 develop-eggs/
 dist/

graphene_sqlalchemy/batching.py

@@ -5,8 +5,13 @@
 import sqlalchemy
 from sqlalchemy.orm import Session, strategies
 from sqlalchemy.orm.query import QueryContext
+from sqlalchemy.util import immutabledict
 
-from .utils import SQL_VERSION_HIGHER_EQUAL_THAN_1_4, is_graphene_version_less_than
+from .utils import (
+    SQL_VERSION_HIGHER_EQUAL_THAN_1_4,
+    SQL_VERSION_HIGHER_EQUAL_THAN_2,
+    is_graphene_version_less_than,
+)
 
 
 def get_data_loader_impl() -> Any:  # pragma: no cover
@@ -76,7 +81,18 @@ async def batch_load_fn(self, parents):
             query_context = parent_mapper_query._compile_context()
         else:
             query_context = QueryContext(session.query(parent_mapper.entity))
-        if SQL_VERSION_HIGHER_EQUAL_THAN_1_4:
+        if SQL_VERSION_HIGHER_EQUAL_THAN_2:  # pragma: no cover
+            self.selectin_loader._load_for_path(
+                query_context,
+                parent_mapper._path_registry,
+                states,
+                None,
+                child_mapper,
+                None,
+                None,  # recursion depth can be none
+                immutabledict(),  # default value for selectinload->lazyload
+            )
+        elif SQL_VERSION_HIGHER_EQUAL_THAN_1_4:
             self.selectin_loader._load_for_path(
                 query_context,
                 parent_mapper._path_registry,

graphene_sqlalchemy/tests/models.py

@@ -16,14 +16,23 @@
     String,
     Table,
     func,
-    select,
 )
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import backref, column_property, composite, mapper, relationship
-from sqlalchemy.sql.sqltypes import _LookupExpressionAdapter
 from sqlalchemy.sql.type_api import TypeEngine
 
+from graphene_sqlalchemy.tests.utils import wrap_select_func
+from graphene_sqlalchemy.utils import SQL_VERSION_HIGHER_EQUAL_THAN_1_4, SQL_VERSION_HIGHER_EQUAL_THAN_2
+
+# fmt: off
+import sqlalchemy
+if SQL_VERSION_HIGHER_EQUAL_THAN_2:
+    from sqlalchemy.sql.sqltypes import HasExpressionLookup # noqa  # isort:skip
+else:
+    from sqlalchemy.sql.sqltypes import _LookupExpressionAdapter as HasExpressionLookup # noqa  # isort:skip
+# fmt: on
+
 PetKind = Enum("cat", "dog", name="pet_kind")
 
 
@@ -119,7 +128,7 @@ def hybrid_prop_list(self) -> List[int]:
         return [1, 2, 3]
 
     column_prop = column_property(
-        select([func.cast(func.count(id), Integer)]), doc="Column property"
+        wrap_select_func(func.cast(func.count(id), Integer)), doc="Column property"
     )
 
     composite_prop = composite(
@@ -163,7 +172,11 @@ def __subclasses__(cls):
 
 editor_table = Table("editors", Base.metadata, autoload=True)
 
-mapper(ReflectedEditor, editor_table)
+# TODO Remove when switching min sqlalchemy version to SQLAlchemy 1.4
+if SQL_VERSION_HIGHER_EQUAL_THAN_1_4:
+    Base.registry.map_imperatively(ReflectedEditor, editor_table)
+else:
+    mapper(ReflectedEditor, editor_table)
 
 
 ############################################
@@ -337,7 +350,7 @@ class Employee(Person):
 ############################################
 
 
-class CustomIntegerColumn(_LookupExpressionAdapter, TypeEngine):
+class CustomIntegerColumn(HasExpressionLookup, TypeEngine):
     """
     Custom Column Type that our converters don't recognize
     Adapted from sqlalchemy.Integer

graphene_sqlalchemy/tests/models_batching.py

@@ -11,11 +11,12 @@
     String,
     Table,
     func,
-    select,
 )
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import column_property, relationship
 
+from graphene_sqlalchemy.tests.utils import wrap_select_func
+
 PetKind = Enum("cat", "dog", name="pet_kind")
 
 
@@ -61,7 +62,7 @@ class Reporter(Base):
     favorite_article = relationship("Article", uselist=False)
 
     column_prop = column_property(
-        select([func.cast(func.count(id), Integer)]), doc="Column property"
+        wrap_select_func(func.cast(func.count(id), Integer)), doc="Column property"
     )
 
 

graphene_sqlalchemy/tests/test_converter.py

@@ -2,20 +2,28 @@
 import sys
 from typing import Dict, Tuple, Union
 
+import graphene
 import pytest
 import sqlalchemy
 import sqlalchemy_utils as sqa_utils
-from sqlalchemy import Column, func, select, types
+from graphene.relay import Node
+from graphene.types.structures import Structure
+from sqlalchemy import Column, func, types
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.inspection import inspect
 from sqlalchemy.orm import column_property, composite
 
-import graphene
-from graphene.relay import Node
-from graphene.types.structures import Structure
-
+from .models import (
+    Article,
+    CompositeFullName,
+    Pet,
+    Reporter,
+    ShoppingCart,
+    ShoppingCartItem,
+)
+from .utils import wrap_select_func
 from ..converter import (
     convert_sqlalchemy_column,
     convert_sqlalchemy_composite,
@@ -27,6 +35,7 @@
 from ..fields import UnsortedSQLAlchemyConnectionField, default_connection_field_factory
 from ..registry import Registry, get_global_registry
 from ..types import ORMField, SQLAlchemyObjectType
+from ..utils import is_sqlalchemy_version_less_than
 from .models import (
     Article,
     CompositeFullName,
@@ -204,9 +213,9 @@ def prop_method() -> int | str:
         return "not allowed in gql schema"
 
     with pytest.raises(
-        ValueError,
-        match=r"Cannot convert hybrid_property Union to "
-        r"graphene.Union: the Union contains scalars. \.*",
+            ValueError,
+            match=r"Cannot convert hybrid_property Union to "
+                  r"graphene.Union: the Union contains scalars. \.*",
     ):
         get_hybrid_property_type(prop_method)
 
@@ -460,7 +469,7 @@ class TestEnum(enum.IntEnum):
 
 def test_should_columproperty_convert():
     field = get_field_from_column(
-        column_property(select([func.sum(func.cast(id, types.Integer))]).where(id == 1))
+        column_property(wrap_select_func(func.sum(func.cast(id, types.Integer))).where(id == 1))
     )
 
     assert field.type == graphene.Int
@@ -477,10 +486,18 @@ def test_should_jsontype_convert_jsonstring():
     assert get_field(types.JSON).type == graphene.JSONString
 
 
+@pytest.mark.skipif(
+    (not is_sqlalchemy_version_less_than("2.0.0b1")),
+    reason="SQLAlchemy >=2.0 does not support this: Variant is no longer used in SQLAlchemy",
+)
 def test_should_variant_int_convert_int():
     assert get_field(types.Variant(types.Integer(), {})).type == graphene.Int
 
 
+@pytest.mark.skipif(
+    (not is_sqlalchemy_version_less_than("2.0.0b1")),
+    reason="SQLAlchemy >=2.0 does not support this: Variant is no longer used in SQLAlchemy",
+)
 def test_should_variant_string_convert_string():
     assert get_field(types.Variant(types.String(), {})).type == graphene.String
 
@@ -811,8 +828,8 @@ class Meta:
     )
 
     for (
-        hybrid_prop_name,
-        hybrid_prop_expected_return_type,
+            hybrid_prop_name,
+            hybrid_prop_expected_return_type,
     ) in shopping_cart_item_expected_types.items():
         hybrid_prop_field = ShoppingCartItemType._meta.fields[hybrid_prop_name]
 
@@ -823,7 +840,7 @@ class Meta:
             str(hybrid_prop_expected_return_type),
         )
         assert (
-            hybrid_prop_field.description is None
+                hybrid_prop_field.description is None
         )  # "doc" is ignored by hybrid property
 
     ###################################################
@@ -870,8 +887,8 @@ class Meta:
     )
 
     for (
-        hybrid_prop_name,
-        hybrid_prop_expected_return_type,
+            hybrid_prop_name,
+            hybrid_prop_expected_return_type,
     ) in shopping_cart_expected_types.items():
         hybrid_prop_field = ShoppingCartType._meta.fields[hybrid_prop_name]
 
@@ -882,5 +899,5 @@ class Meta:
             str(hybrid_prop_expected_return_type),
         )
         assert (
-            hybrid_prop_field.description is None
+                hybrid_prop_field.description is None
         )  # "doc" is ignored by hybrid property

graphene_sqlalchemy/tests/utils.py

@@ -1,6 +1,10 @@
 import inspect
 import re
 
+from sqlalchemy import select
+
+from graphene_sqlalchemy.utils import SQL_VERSION_HIGHER_EQUAL_THAN_1_4
+
 
 def to_std_dicts(value):
     """Convert nested ordered dicts to normal dicts for better comparison."""
@@ -18,8 +22,15 @@ def remove_cache_miss_stat(message):
     return re.sub(r"\[generated in \d+.?\d*s\]\s", "", message)
 
 
-async def eventually_await_session(session, func, *args):
+def wrap_select_func(query):
+    # TODO remove this when we drop support for sqa < 2.0
+    if SQL_VERSION_HIGHER_EQUAL_THAN_1_4:
+        return select(query)
+    else:
+        return select([query])
+
 
+async def eventually_await_session(session, func, *args):
     if inspect.iscoroutinefunction(getattr(session, func)):
         await getattr(session, func)(*args)
     else:

graphene_sqlalchemy/utils.py

@@ -27,12 +27,18 @@ def is_graphene_version_less_than(version_string):  # pragma: no cover
 
 SQL_VERSION_HIGHER_EQUAL_THAN_1_4 = False
 
-if not is_sqlalchemy_version_less_than("1.4"):
+if not is_sqlalchemy_version_less_than("1.4"):  # pragma: no cover
     from sqlalchemy.ext.asyncio import AsyncSession
 
     SQL_VERSION_HIGHER_EQUAL_THAN_1_4 = True
 
 
+SQL_VERSION_HIGHER_EQUAL_THAN_2 = False
+
+if not is_sqlalchemy_version_less_than("2.0.0b1"):  # pragma: no cover
+    SQL_VERSION_HIGHER_EQUAL_THAN_2 = True
+
+
 def get_session(context):
     return context.get("session")
 

setup.py

@@ -15,7 +15,7 @@
     # To keep things simple, we only support newer versions of Graphene
     "graphene>=3.0.0b7",
     "promise>=2.3",
-    "SQLAlchemy>=1.1,<2",
+    "SQLAlchemy>=1.1",
     "aiodataloader>=0.2.0,<1.0",
 ]
 

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = pre-commit,py{37,38,39,310}-sql{12,13,14}
+envlist = pre-commit,py{37,38,39,310}-sql{12,13,14,20}
 skipsdist = true
 minversion = 3.7.0
 
@@ -15,6 +15,7 @@ SQLALCHEMY =
     1.2: sql12
     1.3: sql13
     1.4: sql14
+    2.0: sql20
 
 [testenv]
 passenv = GITHUB_*
@@ -23,8 +24,11 @@ deps =
     sql12: sqlalchemy>=1.2,<1.3
     sql13: sqlalchemy>=1.3,<1.4
     sql14: sqlalchemy>=1.4,<1.5
+    sql20: sqlalchemy>=2.0.0b3
+setenv =
+    SQLALCHEMY_WARN_20 = 1
 commands =
-    pytest graphene_sqlalchemy --cov=graphene_sqlalchemy --cov-report=term --cov-report=xml {posargs}
+    python -W always -m pytest graphene_sqlalchemy --cov=graphene_sqlalchemy --cov-report=term --cov-report=xml {posargs}
 
 [testenv:pre-commit]
 basepython=python3.10