Introduce private sqlalchemy 2.0 compatibility helpers

Author: kurtmckee

conftest.py

@@ -3,8 +3,9 @@
 
 import pytest
 import sqlalchemy as sa
+import sqlalchemy.event
+import sqlalchemy.exc
 from sqlalchemy import create_engine
-from sqlalchemy.ext.declarative import declarative_base, synonym_for
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.orm.session import close_all_sessions
@@ -15,6 +16,11 @@
     i18n,
     InstrumentedList
 )
+from sqlalchemy_utils.compat import (
+    _declarative_base,
+    _select_args,
+    _synonym_for
+)
 from sqlalchemy_utils.functions.orm import _get_class_registry
 from sqlalchemy_utils.types.pg_composite import remove_composite_listeners
 
@@ -148,7 +154,7 @@ def connection(engine):
 
 @pytest.fixture
 def Base():
-    return declarative_base()
+    return _declarative_base()
 
 
 @pytest.fixture
@@ -185,7 +191,7 @@ def articles_count(self):
         def articles_count(cls):
             Article = _get_class_registry(Base)['Article']
             return (
-                sa.select([sa.func.count(Article.id)])
+                sa.select(*_select_args(sa.func.count(Article.id)))
                 .where(Article.category_id == cls.id)
                 .correlate(Article.__table__)
                 .label('article_count')
@@ -195,7 +201,7 @@ def articles_count(cls):
         def name_alias(self):
             return self.name
 
-        @synonym_for('name')
+        @_synonym_for('name')
         @property
         def name_synonym(self):
             return self.name
@@ -229,15 +235,22 @@ def init_models(User, Category, Article):
 @pytest.fixture
 def session(request, engine, connection, Base, init_models):
     sa.orm.configure_mappers()
-    Base.metadata.create_all(connection)
+    with connection.begin():
+        Base.metadata.create_all(connection)
     Session = sessionmaker(bind=connection)
-    session = Session()
+    try:
+        # Enable sqlalchemy 2.0 behavior.
+        session = Session(future=True)
+    except TypeError:
+        # sqlalchemy 1.3
+        session = Session()
     i18n.get_locale = get_locale
 
     def teardown():
         aggregates.manager.reset()
         close_all_sessions()
-        Base.metadata.drop_all(connection)
+        with connection.begin():
+            Base.metadata.drop_all(connection)
         remove_composite_listeners()
         connection.close()
         engine.dispose()

setup.py

@@ -75,6 +75,7 @@ def get_version():
     platforms='any',
     install_requires=[
         'SQLAlchemy>=1.0',
+        "importlib_metadata ; python_version<'3.8'",
     ],
     extras_require=extras_require,
     python_requires='~=3.6',

sqlalchemy_utils/compat.py

@@ -1,5 +1,75 @@
-def get_scalar_subquery(query):
-    try:
+import sys
+
+if sys.version_info >= (3, 8):
+    from importlib.metadata import metadata
+else:
+    from importlib_metadata import metadata
+
+
+_sqlalchemy_version = tuple(
+    [int(i) for i in metadata("sqlalchemy")["Version"].split(".")[:2]]
+)
+
+
+# In sqlalchemy 2.0, some functions moved to sqlalchemy.orm.
+# In sqlalchemy 1.3, they are only available in .ext.declarative.
+# In sqlalchemy 1.4, they are available in both places.
+#
+# WARNING
+# -------
+#
+# These imports are for internal, private compatibility.
+# They are not supported and may change or move at any time.
+# Do not import these in your own code.
+#
+
+if _sqlalchemy_version >= (1, 4):
+    from sqlalchemy.orm import declarative_base as _declarative_base
+    from sqlalchemy.orm import synonym_for as _synonym_for
+else:
+    from sqlalchemy.ext.declarative import declarative_base as _declarative_base
+    from sqlalchemy.ext.declarative import synonym_for as _synonym_for
+
+
+# scalar subqueries
+if _sqlalchemy_version >= (1, 4):
+    def get_scalar_subquery(query):
         return query.scalar_subquery()
-    except AttributeError:  # SQLAlchemy <1.4
+else:
+    def get_scalar_subquery(query):
         return query.as_scalar()
+
+
+# In sqlalchemy 2.0, select() columns are positional.
+# In sqlalchemy 1.3, select() columns must be wrapped in a list.
+#
+# _select_args() is designed so its return value can be unpacked:
+#
+#     select(*_select_args(1, 2))
+#
+# When sqlalchemy 1.3 support is dropped, remove the call to _select_args()
+# and keep the arguments the same:
+#
+#     select(1, 2)
+#
+# WARNING
+# -------
+#
+# _select_args() is a private, internal function.
+# It is not supported and may change or move at any time.
+# Do not import this in your own code.
+#
+if _sqlalchemy_version >= (1, 4):
+    def _select_args(*args):
+        return args
+else:
+    def _select_args(*args):
+        return [args]
+
+
+__all__ = (
+    "_declarative_base",
+    "get_scalar_subquery",
+    "_select_args",
+    "_synonym_for",
+)