support sqlalchemy from 1.4.18 to 2.0+, fix #12

h0rn3t · h0rn3t · commit b2c4be449584 · 2023-08-16T08:41:07.000+03:00
diff --git a/README.md b/README.md
@@ -20,7 +20,8 @@ Provides SQLAlchemy middleware for FastAPI using AsyncSession and async engine.
 
 ### Important !!!
 If you use ```sqlmodel``` install ```fastapi-async-sqlalchemy==0.3.14```, because it depends on ```sqlalchemy<=1.4.41```
-```bash 
+```last version``` compatible with ```sqlmodel``` sqlalchemy>=2.0.0 but if you will use ```sqlalchemy>=1.4.42``` and ```sqlmodel``` you will get error
+```bash
 pip install fastapi-async-sqlalchemy==0.3.14
  ```
 
diff --git a/fastapi_async_sqlalchemy/__init__.py b/fastapi_async_sqlalchemy/__init__.py
@@ -2,4 +2,4 @@
 
 __all__ = ["db", "SQLAlchemyMiddleware"]
 
-__version__ = "0.4.2"
+__version__ = "0.5.0"
diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py
@@ -3,7 +3,13 @@
 
 from sqlalchemy.engine import Engine
 from sqlalchemy.engine.url import URL
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+
+try:
+    from sqlalchemy.ext.asyncio import async_sessionmaker
+except ImportError:
+    from sqlalchemy.orm import sessionmaker as async_sessionmaker
+
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from starlette.requests import Request
 from starlette.types import ASGIApp
@@ -37,16 +43,16 @@ def __init__(
             engine = custom_engine
 
         global _Session
-        _Session = async_sessionmaker(engine, expire_on_commit=False, **session_args)
+        _Session = async_sessionmaker(
+            engine, class_=AsyncSession, expire_on_commit=False, **session_args
+        )
 
     async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
         async with db(commit_on_exit=self.commit_on_exit):
             return await call_next(request)
 
 
 class DBSessionMeta(type):
-    # using this metaclass means that we can access db.session as a property at a class level,
-    # rather than db().session
     @property
     def session(self) -> AsyncSession:
         """Return an instance of Session local to the current async context."""
@@ -75,14 +81,15 @@ async def __aenter__(self):
 
     async def __aexit__(self, exc_type, exc_value, traceback):
         session = _session.get()
-        if exc_type is not None:
-            await session.rollback()
-
-        if self.commit_on_exit:
-            await session.commit()
 
-        await session.close()
-        _session.reset(self.token)
+        try:
+            if exc_type is not None:
+                await session.rollback()
+            elif self.commit_on_exit:  # Note: Changed this to elif to avoid commit after rollback
+                await session.commit()
+        finally:
+            await session.close()
+            _session.reset(self.token)
 
 
 db: DBSessionMeta = DBSession
diff --git a/requirements.txt b/requirements.txt
@@ -27,7 +27,7 @@ regex==2020.2.20
 requests>=2.22.0
 httpx>=0.20.0
 six==1.12.0
-SQLAlchemy>=2.0.0
+SQLAlchemy>=1.4.19
 asyncpg>=0.27.0
 aiosqlite==0.19.0
 sqlparse==0.4.4
diff --git a/setup.py b/setup.py
@@ -27,7 +27,7 @@
     package_data={"fastapi_async_sqlalchemy": ["py.typed"]},
     zip_safe=False,
     python_requires=">=3.7",
-    install_requires=["starlette>=0.13.6", "SQLAlchemy>=2.0.0"],
+    install_requires=["starlette>=0.13.6", "SQLAlchemy>=1.4.19"],
     classifiers=[
         "Development Status :: 4 - Beta",
         "Environment :: Web Environment",
diff --git a/tests/test_session.py b/tests/test_session.py
@@ -2,6 +2,7 @@
 import sys
 
 import pytest
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 from starlette.middleware.base import BaseHTTPMiddleware
 
@@ -119,31 +120,14 @@ async def test_init_session(app, db, SQLAlchemyMiddleware):
         assert isinstance(db.session, AsyncSession)
 
 
-@pytest.mark.parametrize("commit_on_exit", [True, False])
 @pytest.mark.asyncio
-async def test_db_context_session_args(app, db, SQLAlchemyMiddleware, commit_on_exit):
-    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url, commit_on_exit=commit_on_exit)
+async def test_db_session_commit_fail(app, db, SQLAlchemyMiddleware):
+    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url, commit_on_exit=True)
 
-    session_args = {}
-
-    async with db(session_args=session_args, commit_on_exit=True):
-        assert isinstance(db.session, AsyncSession)
-
-        # assert db.session.expire_on_commit
-
-    session_args = {"expire_on_commit": False}
-    async with db(session_args=session_args):
-        db.session
-        # assert db.session.expire_on_commit
-
-
-@pytest.mark.asyncio
-async def test_rollback(app, db, SQLAlchemyMiddleware):
-    #  pytest-cov shows that the line in db.__exit__() rolling back the db session
-    #  when there is an Exception is run correctly. However, it would be much better
-    #  if we could demonstrate somehow that db.session.rollback() was called e.g. once
-    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url)
-
-    with pytest.raises(Exception):
+    with pytest.raises(IntegrityError):
         async with db():
-            raise Exception
+            raise IntegrityError("test", "test", "test")
+        db.session.close.assert_called_once()
+
+    async with db():
+        assert db.session

Original file line number	Diff line number	Diff line change
`@@ -2,4 +2,4 @@`
`2`	`2`
`3`	`3`	`__all__ = ["db", "SQLAlchemyMiddleware"]`
`4`	`4`
`5`		`-__version__ = "0.4.2"`
	`5`	`+__version__ = "0.5.0"`