SQLAlchemy 2.0 compatibility (#901)

* SQLAlchemy 2.0 support

* codestyle

* SQLAlchemy 2.0 style examples

* Run tests sqlalchemy v1 and v2

* Bump min supported sqlalchemy version

---------

Co-authored-by: Iurii Pliner <yury.pliner@gmail.com>

.github/workflows/ci.yml

@@ -17,6 +17,7 @@ jobs:
     strategy:
       matrix:
         python: ['3.7', '3.8', '3.9', '3.10', '3.11']
+        sqlalchemy: ['sqlalchemy[postgresql_psycopg2binary]>=1.4,<1.5', 'sqlalchemy[postgresql_psycopg2binary]>=2.0,<2.1']
       fail-fast: false
     runs-on: ubuntu-latest
     timeout-minutes: 20
@@ -32,6 +33,9 @@ jobs:
         pip install -e .
         pip install -r requirements.txt
         pip install codecov
+    - name: Install ${{ matrix.sqlalchemy }}
+      run: |
+        pip install "${{ matrix.sqlalchemy }}"
     - name: Run tests
       run: |
         make cov-ci

aiopg/sa/engine.py

@@ -8,10 +8,19 @@
 from .connection import SAConnection
 
 try:
-    from sqlalchemy.dialects.postgresql.psycopg2 import (
-        PGCompiler_psycopg2,
-        PGDialect_psycopg2,
-    )
+    from sqlalchemy import __version__
+
+    sa_version = tuple(map(int, __version__.split(".")))
+    if sa_version[0] < 2:
+        from sqlalchemy.dialects.postgresql.psycopg2 import (
+            PGCompiler_psycopg2,
+            PGDialect_psycopg2,
+        )
+    else:
+        from sqlalchemy.dialects.postgresql.base import (
+            PGCompiler as PGCompiler_psycopg2,
+        )
+        from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
 except ImportError:  # pragma: no cover
     raise ImportError("aiopg.sa requires sqlalchemy")
 

aiopg/sa/result.py

@@ -108,15 +108,21 @@ def __init__(self, result_proxy, cursor_description):
         # older versions of SQLAlchemy.
         typemap = getattr(dialect, "dbapi_type_map", {})
 
-        assert (
-            dialect.case_sensitive
+        # `case_sensitive` property removed in SQLAlchemy 2.0+.
+        # Usage of `getattr` only needed for backward compatibility with
+        # older versions of SQLAlchemy.
+        assert getattr(
+            dialect, "case_sensitive", True
         ), "Doesn't support case insensitive database connection"
 
         # high precedence key values.
         primary_keymap = {}
 
-        assert (
-            not dialect.description_encoding
+        # `description_encoding` property removed in SQLAlchemy 2.0+.
+        # Usage of `getattr` only needed for backward compatibility with
+        # older versions of SQLAlchemy.
+        assert not getattr(
+            dialect, "description_encoding", None
         ), "psycopg in py3k should not use this"
 
         for i, rec in enumerate(cursor_description):

docs/index.rst

@@ -32,7 +32,7 @@ Current version is |release|.
 .. warning::
     1. Removing await the before :meth:`Cursor.mogrify` function
 
-    2. Only supports ``python >= 3.6``
+    2. Only supports ``python >= 3.7``
 
     3. Only support syntax ``async/await``
 

examples/isolation_sa_transaction.py

@@ -3,7 +3,7 @@
 import sqlalchemy as sa
 from psycopg2 import InternalError
 from psycopg2.extensions import TransactionRollbackError
-from sqlalchemy.sql.ddl import CreateTable
+from sqlalchemy.sql.ddl import CreateTable, DropTable
 
 from aiopg.sa import create_engine
 
@@ -18,6 +18,7 @@
 
 
 async def create_sa_transaction_tables(conn):
+    await conn.execute(DropTable(users, if_exists=True))
     await conn.execute(CreateTable(users))
 
 

examples/sa.py

@@ -100,12 +100,12 @@ async def fill_data(conn):
 
 
 async def count(conn):
-    c1 = await conn.scalar(users.count())
-    c2 = await conn.scalar(emails.count())
+    c1 = await conn.scalar(sa.select(sa.func.count(users.c.id)))
+    c2 = await conn.scalar(sa.select(sa.func.count(emails.c.id)))
     print("Population consists of", c1, "people with", c2, "emails in total")
     join = sa.join(emails, users, users.c.id == emails.c.user_id)
     query = (
-        sa.select([users.c.name])
+        sa.select(users.c.name)
         .select_from(join)
         .where(emails.c.private == False)  # noqa
         .group_by(users.c.name)
@@ -121,11 +121,11 @@ async def count(conn):
 
 async def show_julia(conn):
     print("Lookup for Julia:")
-    join = sa.join(emails, users, users.c.id == emails.c.user_id)
     query = (
-        sa.select([users, emails], use_labels=True)
-        .select_from(join)
+        sa.select(users, emails)
+        .join(emails, users.c.id == emails.c.user_id)
         .where(users.c.name == "Julia")
+        .set_label_style(sa.LABEL_STYLE_TABLENAME_PLUS_COL)
     )
     async for row in conn.execute(query):
         print(
@@ -138,9 +138,7 @@ async def show_julia(conn):
 
 
 async def ave_age(conn):
-    query = sa.select(
-        [sa.func.avg(sa.func.age(users.c.birthday))]
-    ).select_from(users)
+    query = sa.select(sa.func.avg(sa.func.age(users.c.birthday)))
     ave = await conn.scalar(query)
     print(
         "Average age of population is",

examples/simple_sa_transaction.py

@@ -1,7 +1,7 @@
 import asyncio
 
 import sqlalchemy as sa
-from sqlalchemy.schema import CreateTable
+from sqlalchemy.schema import CreateTable, DropTable
 
 from aiopg.sa import create_engine
 
@@ -16,6 +16,7 @@
 
 
 async def create_sa_transaction_tables(conn):
+    await conn.execute(DropTable(users, if_exists=True))
     await conn.execute(CreateTable(users))
 
 

examples/transaction.py

@@ -21,6 +21,7 @@ async def main():
     async with create_pool(dsn) as pool:
         async with pool.acquire() as conn:
             async with conn.cursor() as cur:
+                await cur.execute("DROP TABLE IF EXISTS tbl")
                 await cur.execute("CREATE TABLE tbl (id int)")
                 await transaction(cur, IsolationLevel.repeatable_read)
                 await transaction(cur, IsolationLevel.read_committed)

requirements.txt

@@ -1,5 +1,5 @@
 coverage==5.5
-docker==5.0.0
+docker==6.1.2
 flake8==3.9.2
 isort==5.9.3
 -e .[sa]
@@ -10,7 +10,7 @@ pytest-sugar==0.9.4
 pytest-timeout==1.4.2
 sphinxcontrib-asyncio==0.3.0
 psycopg2-binary==2.9.5
-sqlalchemy[postgresql_psycopg2binary]==1.4.38
+sqlalchemy[postgresql_psycopg2binary]==2.0.20
 async-timeout==4.0.0
 mypy==0.910
 black==22.3.0

setup.py

@@ -4,7 +4,7 @@
 from setuptools import setup, find_packages
 
 install_requires = ["psycopg2-binary>=2.9.5", "async_timeout>=3.0,<5.0"]
-extras_require = {"sa": ["sqlalchemy[postgresql_psycopg2binary]>=1.3,<1.5"]}
+extras_require = {"sa": ["sqlalchemy[postgresql_psycopg2binary]>=1.4,<2.1"]}
 
 
 def read(*parts):

tests/test_sa_connection.py

@@ -133,7 +133,7 @@ async def test_execute_sa_insert_positional_params(connect):
 
 async def test_scalar(connect):
     conn = await connect()
-    res = await conn.scalar(select([func.count()]).select_from(tbl))
+    res = await conn.scalar(select(func.count()).select_from(tbl))
     assert 1, res
 
 

tests/test_sa_priority_name.py

@@ -1,5 +1,6 @@
 import pytest
 import sqlalchemy as sa
+from sqlalchemy import LABEL_STYLE_TABLENAME_PLUS_COL
 
 meta = sa.MetaData()
 tbl = sa.Table(
@@ -36,7 +37,7 @@ async def test_priority_name(connect):
 
 async def test_priority_name_label(connect):
     await connect.execute(tbl.insert().values(id="test_id", name="test_name"))
-    query = sa.select([tbl.c.name.label("test_label_name"), tbl.c.id])
+    query = sa.select(tbl.c.name.label("test_label_name"), tbl.c.id)
     query = query.select_from(tbl)
     row = await (await connect.execute(query)).first()
     assert row.test_label_name == "test_name"
@@ -46,7 +47,7 @@ async def test_priority_name_label(connect):
 async def test_priority_name_and_label(connect):
     await connect.execute(tbl.insert().values(id="test_id", name="test_name"))
     query = sa.select(
-        [tbl.c.name.label("test_label_name"), tbl.c.name, tbl.c.id]
+        tbl.c.name.label("test_label_name"), tbl.c.name, tbl.c.id
     )
     query = query.select_from(tbl)
     row = await (await connect.execute(query)).first()
@@ -57,7 +58,7 @@ async def test_priority_name_and_label(connect):
 
 async def test_priority_name_all_get(connect):
     await connect.execute(tbl.insert().values(id="test_id", name="test_name"))
-    query = sa.select([tbl.c.name])
+    query = sa.select(tbl.c.name)
     query = query.select_from(tbl)
     row = await (await connect.execute(query)).first()
     assert row.name == "test_name"
@@ -69,7 +70,7 @@ async def test_priority_name_all_get(connect):
 async def test_use_labels(connect):
     """key property is ignored"""
     await connect.execute(tbl.insert().values(id="test_id", name="test_name"))
-    query = tbl.select(use_labels=True)
+    query = tbl.select().set_label_style(LABEL_STYLE_TABLENAME_PLUS_COL)
     row = await (await connect.execute(query)).first()
     assert row.sa_tbl5_Name == "test_name"
     assert row.sa_tbl5_ID == "test_id"