chore: prepare fastapi-sqla to be compliant with sqlalchemy 2.0 - DIA-40134 (#48)

Author: hadrien

fastapi_sqla/__init__.py

@@ -11,7 +11,7 @@
 from fastapi.responses import PlainTextResponse
 from pydantic import BaseModel, Field
 from pydantic.generics import GenericModel
-from sqlalchemy import engine_from_config
+from sqlalchemy import engine_from_config, text
 from sqlalchemy.ext.declarative import DeferredReflection
 from sqlalchemy.orm import Query as LegacyQuery
 from sqlalchemy.orm.session import Session as SqlaSession
@@ -63,7 +63,9 @@ def setup(app: FastAPI):
 
 
 def startup():
-    lowercase_environ = {k.lower(): v for k, v in os.environ.items()}
+    lowercase_environ = {
+        k.lower(): v for k, v in os.environ.items() if k.lower() != "sqlalchemy_warn_20"
+    }
     engine = engine_from_config(lowercase_environ, prefix="sqlalchemy_")
     aws_rds_iam_support.setup(engine.engine)
 
@@ -74,7 +76,7 @@ def startup():
     # Fail early:
     try:
         with open_session() as session:
-            session.execute("select 'OK'")
+            session.execute(text("select 'OK'"))
     except Exception:
         logger.critical(
             "Fail querying db: is sqlalchemy_url envvar correctly configured?"
@@ -234,7 +236,9 @@ def default_query_count(session: Session, query: DbQuery) -> int:
         result = query.count()
 
     elif isinstance(query, Select):
-        result = session.execute(select(func.count()).select_from(query)).scalar()
+        result = session.execute(
+            select(func.count()).select_from(query.subquery())
+        ).scalar()
 
     else:  # pragma no cover
         raise NotImplementedError(f"Query type {type(query)!r} is not supported")

fastapi_sqla/_pytest_plugin.py

@@ -16,11 +16,19 @@
     asyncio_support = False
 
 
+@fixture(scope="session")
+def sqla_version_tuple():
+    from sqlalchemy import __version__
+
+    return tuple(int(i) for i in __version__.split("."))
+
+
 @fixture(scope="session")
 def db_host():
     """Default db host used by depending fixtures.
 
-    When CI key is set in environment variables, it uses `postgres` as host name else, host used is `localhost`
+    When CI key is set in environment variables, it uses `postgres` as host name else,
+    host used is `localhost`
     """
 
     return "postgres" if "CI" in os.environ else "localhost"
@@ -49,9 +57,8 @@ def db_url(db_host, db_user):
 @fixture(scope="session")
 def sqla_connection(db_url):
     engine = create_engine(db_url)
-    connection = engine.connect()
-    yield connection
-    connection.close()
+    with engine.connect() as connection:
+        yield connection
 
 
 @fixture(scope="session")
@@ -69,7 +76,10 @@ def db_migration(db_url, sqla_connection, alembic_ini_path):
     alembic_config = Config(file_=alembic_ini_path)
     alembic_config.set_main_option("sqlalchemy.url", db_url)
 
-    sqla_connection.execute(text("DROP SCHEMA public CASCADE; CREATE SCHEMA public;"))
+    with sqla_connection.begin():
+        sqla_connection.execute(
+            text("DROP SCHEMA public CASCADE; CREATE SCHEMA public;")
+        )
 
     command.upgrade(alembic_config, "head")
     yield
@@ -112,7 +122,7 @@ def sqla_transaction(sqla_connection):
 
 
 @fixture
-def session(sqla_transaction, sqla_connection):
+def session(sqla_transaction, sqla_connection, sqla_version_tuple):
     """Sqla session to use when creating db fixtures.
 
     While it does not write any record in DB, the application will still be able to
@@ -121,8 +131,15 @@ def session(sqla_transaction, sqla_connection):
     import fastapi_sqla
 
     session = fastapi_sqla._Session(bind=sqla_connection)
-    yield session
-    session.close()
+
+    if sqla_version_tuple >= (1, 4, 0):
+        with session.begin():
+            yield session
+            session.rollback()
+    else:
+        with session.begin(subtransactions=True):
+            session.begin_nested()
+            yield session
 
 
 def format_async_async_sqlalchemy_url(url):

pyproject.toml

@@ -110,3 +110,7 @@ extras =
     tests
 commands = pytest -vv --cov={envsitepackagesdir}/fastapi_sqla --cov-report xml --cov-report html --junitxml=test-reports/pytest/junit.xml
 """
+
+[tool.black]
+exclude = ".mypy_cache|.pytest_cache|.vscode|.eggs|venv"
+--skip-magic-trailing-comma = true

setup.cfg

@@ -8,6 +8,9 @@ addopts =
     --cov-report term
     --cov-report term-missing
 
+filterwarnings =
+    error:.*removed in version 2\.0.*:
+
 [pytest-watch]
 ext = .py,.yaml,.cfg,.yml
 

tests/conftest.py

@@ -1,10 +1,12 @@
 import importlib
+import os
 from unittest.mock import patch
 
 from faker import Faker
 from pytest import fixture, skip
-from sqlalchemy import engine_from_config
-from sqlalchemy.orm.session import close_all_sessions
+
+# Must be done before importing anything from sqlalchemy:
+os.environ["SQLALCHEMY_WARN_20"] = "true"
 
 pytest_plugins = ["fastapi_sqla._pytest_plugin", "pytester"]
 
@@ -23,6 +25,7 @@ def pytest_configure(config):
     config.addinivalue_line(
         "markers", "require_boto3: skip test if boto3 is not installed"
     )
+    config.addinivalue_line("markers", "dont_patch_engines: do not patch engines")
 
 
 @fixture(scope="session")
@@ -52,7 +55,11 @@ def is_boto3_installed():
 
 @fixture(scope="session", autouse=True)
 def environ(db_url, sqla_version_tuple, async_sqlalchemy_url):
-    values = {"sqlalchemy_url": db_url, "PYTHONASYNCIODEBUG": "1"}
+    values = {
+        "PYTHONASYNCIODEBUG": "1",
+        "sqlalchemy_url": db_url,
+        "SQLALCHEMY_WARN_20": "true",
+    }
 
     if sqla_version_tuple >= (1, 4, 0) and is_asyncpg_installed():
         values["async_sqlalchemy_url"] = async_sqlalchemy_url
@@ -63,12 +70,16 @@ def environ(db_url, sqla_version_tuple, async_sqlalchemy_url):
 
 @fixture(scope="session")
 def engine(environ):
+    from sqlalchemy import engine_from_config
+
     engine = engine_from_config(environ, prefix="sqlalchemy_")
     return engine
 
 
 @fixture(autouse=True)
-def tear_down():
+def tear_down(environ):
+    from sqlalchemy.orm.session import close_all_sessions
+
     import fastapi_sqla
 
     yield

tests/test_base.py

@@ -6,11 +6,13 @@
 @fixture(autouse=True, scope="module")
 def setup_tear_down(engine):
     with engine.connect() as connection:
-        connection.execute(
-            text("CREATE TABLE IF NOT EXISTS test_table (id integer primary key)")
-        )
+        with connection.begin():
+            connection.execute(
+                text("CREATE TABLE IF NOT EXISTS test_table (id integer primary key)")
+            )
         yield
-        connection.execute(text("DROP TABLE test_table"))
+        with connection.begin():
+            connection.execute(text("DROP TABLE test_table"))
 
 
 def test_startup_reflect_test_table():

tests/test_middleware.py

@@ -14,19 +14,21 @@
 @fixture(scope="module", autouse=True)
 def setup_tear_down(engine):
     with engine.connect() as connection:
-        connection.execute(
-            text(
-                """
-                CREATE TABLE IF NOT EXISTS public.user (
-                   id integer primary key,
-                   first_name varchar,
-                   last_name varchar
+        with connection.begin():
+            connection.execute(
+                text(
+                    """
+                    CREATE TABLE IF NOT EXISTS public.user (
+                       id integer primary key,
+                       first_name varchar,
+                       last_name varchar
+                    )
+                    """
                 )
-                """
             )
-        )
         yield
-        connection.execute(text("DROP TABLE public.user"))
+        with connection.begin():
+            connection.execute(text("DROP TABLE public.user"))
 
 
 @fixture
@@ -104,11 +106,12 @@ async def test_session_dependency(client, faker, session):
     first_name = faker.first_name()
     last_name = faker.last_name()
     res = await client.post(
-        "/users",
-        json={"id": userid, "first_name": first_name, "last_name": last_name},
+        "/users", json={"id": userid, "first_name": first_name, "last_name": last_name}
     )
     assert res.status_code == 200, res.json()
-    row = session.execute(f"select * from public.user where id = {userid}").fetchone()
+    row = session.execute(
+        text(f"select * from public.user where id = {userid}")
+    ).fetchone()
     assert row == (userid, first_name, last_name)
 
 
@@ -124,7 +127,9 @@ async def test_async_session_dependency(client, faker, async_session):
     )
     assert res.status_code == 200, res.json()
     row = (
-        await async_session.execute(f"select * from public.user where id = {userid}")
+        await async_session.execute(
+            text(f"select * from public.user where id = {userid}")
+        )
     ).fetchone()
     assert row == (userid, first_name, last_name)
 

tests/test_open_session.py

@@ -5,11 +5,16 @@
 
 @fixture(autouse=True, scope="module")
 def module_setup_tear_down(engine, sqla_connection):
-    engine.execute(
-        "CREATE TABLE IF NOT EXISTS test_table (id integer primary key, value varchar)"
-    )
+    with sqla_connection.begin():
+        sqla_connection.execute(
+            text(
+                "CREATE TABLE IF NOT EXISTS test_table   "
+                "(id integer primary key, value varchar) "
+            )
+        )
     yield
-    engine.execute("DROP TABLE test_table")
+    with sqla_connection.begin():
+        sqla_connection.execute(text("DROP TABLE test_table"))
 
 
 @fixture(autouse=True)

tests/test_pagination.py

@@ -13,37 +13,39 @@
 @fixture(scope="module", autouse=True)
 def setup_tear_down(sqla_connection):
     faker = Faker(seed=0)
-    sqla_connection.execute(
-        text(
-            "create table if not exists public.user "
-            "(id serial primary key, name varchar)"
+    with sqla_connection.begin():
+        sqla_connection.execute(
+            text(
+                "create table if not exists public.user "
+                "(id serial primary key, name varchar)"
+            )
         )
-    )
-    sqla_connection.execute(
-        text(
+        sqla_connection.execute(
+            text(
+                """
+            create table if not exists note (
+                user_id integer,
+                id serial,
+                content text,
+                primary key (user_id, id),
+                foreign key (user_id) references public.user (id)
+            )
             """
-        create table if not exists note (
-            user_id integer,
-            id serial,
-            content text,
-            primary key (user_id, id),
-            foreign key (user_id) references public.user (id)
-        )
-        """
+            )
         )
-    )
-    metadata = MetaData()
-    user = Table("user", metadata, autoload_with=sqla_connection)
-    note = Table("note", metadata, autoload_with=sqla_connection)
-    user_params = [{"name": faker.name()} for _ in range(1, 43)]
-    note_params = [
-        {"user_id": i % 42 + 1, "content": faker.text()} for i in range(0, 22 * 42)
-    ]
-    sqla_connection.execute(user.insert(), *user_params)
-    sqla_connection.execute(note.insert(), *note_params)
+        metadata = MetaData()
+        user = Table("user", metadata, autoload_with=sqla_connection)
+        note = Table("note", metadata, autoload_with=sqla_connection)
+        user_params = [{"name": faker.name()} for _ in range(1, 43)]
+        note_params = [
+            {"user_id": i % 42 + 1, "content": faker.text()} for i in range(0, 22 * 42)
+        ]
+        sqla_connection.execute(user.insert(), user_params)
+        sqla_connection.execute(note.insert(), note_params)
     yield
-    sqla_connection.execute(text("drop table note cascade"))
-    sqla_connection.execute(text("drop table public.user cascade"))
+    with sqla_connection.begin():
+        sqla_connection.execute(text("drop table note cascade"))
+        sqla_connection.execute(text("drop table public.user cascade"))
 
 
 @fixture
@@ -80,7 +82,7 @@ class Note(Base):
 def test_pagination(session, user_cls, offset, limit, total_pages, page_number):
     from fastapi_sqla import Paginate
 
-    query = session.query(user_cls).options(joinedload("notes"))
+    query = session.query(user_cls).options(joinedload(user_cls.notes))
     result = Paginate(session, offset, limit)(query)
 
     assert result.meta.total_items == 42
@@ -99,7 +101,7 @@ def test_pagination_with_legacy_query_count(
 ):
     from fastapi_sqla import Paginate
 
-    query = session.query(user_cls).options(joinedload("notes"))
+    query = session.query(user_cls).options(joinedload(user_cls.notes))
     result = Paginate(session, offset, limit)(query)
 
     assert result.meta.total_items == 42
@@ -154,13 +156,17 @@ class UserWithMeta(User):
     @app.get("/v1/users", response_model=Page[UserWithNotes])
     def sqla_13_all_users(session: Session = Depends(), paginate: Paginate = Depends()):
         query = (
-            session.query(user_cls).options(joinedload("notes")).order_by(user_cls.id)
+            session.query(user_cls)
+            .options(joinedload(user_cls.notes))
+            .order_by(user_cls.id)
         )
         return paginate(query)
 
     @app.get("/v2/users", response_model=Page[UserWithNotes])
     def sqla_14_all_users(paginate: Paginate = Depends()):
-        query = select(user_cls).options(joinedload("notes")).order_by(user_cls.id)
+        query = (
+            select(user_cls).options(joinedload(user_cls.notes)).order_by(user_cls.id)
+        )
         return paginate(query)
 
     @app.get("/v2/users-with-notes-count", response_model=Page[UserWithNotesCount])
@@ -182,8 +188,7 @@ def query_with_JSON_result(paginate: Paginate = Depends()):
                 user_cls.id,
                 user_cls.name,
                 cast(
-                    func.format('{"notes_count": %s}', func.count(note_cls.id)),
-                    JSON,
+                    func.format('{"notes_count": %s}', func.count(note_cls.id)), JSON
                 ).label("meta"),
             )
             .join(note_cls)