fastapi-async-sqlalchemy 0.6.0__tar.gz → 0.7.0.dev1__tar.gz

{fastapi-async-sqlalchemy-0.6.0 → fastapi_async_sqlalchemy-0.7.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastapi-async-sqlalchemy
-Version: 0.6.0
+Version: 0.7.0.dev1
 Summary: SQLAlchemy middleware for FastAPI
 Home-page: https://github.com/h0rn3t/fastapi-async-sqlalchemy.git
 Author: Eugene Shershen
@@ -8,7 +8,7 @@ Author-email: h0rn3t.null@gmail.com
 License: MIT
 Project-URL: Code, https://github.com/h0rn3t/fastapi-async-sqlalchemy
 Project-URL: Issue tracker, https://github.com/h0rn3t/fastapi-async-sqlalchemy/issues
-Classifier: Development Status :: 4 - Beta
+Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: Web Environment
 Classifier: Framework :: AsyncIO
 Classifier: Intended Audience :: Developers
@@ -19,6 +19,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Topic :: Internet :: WWW/HTTP :: HTTP Servers
@@ -27,6 +28,8 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.7
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: starlette>=0.13.6
+Requires-Dist: SQLAlchemy>=1.4.19
 
 # SQLAlchemy FastAPI middleware
 
@@ -35,7 +38,7 @@ License-File: LICENSE
 [![codecov](https://codecov.io/gh/h0rn3t/fastapi-async-sqlalchemy/branch/main/graph/badge.svg?token=F4NJ34WKPY)](https://codecov.io/gh/h0rn3t/fastapi-async-sqlalchemy)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 [![pip](https://img.shields.io/pypi/v/fastapi_async_sqlalchemy?color=blue)](https://pypi.org/project/fastapi-async-sqlalchemy/)
-[![Downloads](https://pepy.tech/badge/fastapi-async-sqlalchemy)](https://pepy.tech/project/fastapi-async-sqlalchemy)
+[![Downloads](https://static.pepy.tech/badge/fastapi-async-sqlalchemy)](https://pepy.tech/project/fastapi-async-sqlalchemy)
 [![Updates](https://pyup.io/repos/github/h0rn3t/fastapi-async-sqlalchemy/shield.svg)](https://pyup.io/repos/github/h0rn3t/fastapi-async-sqlalchemy/)
 
 ### Description
@@ -48,8 +51,8 @@ Provides SQLAlchemy middleware for FastAPI using AsyncSession and async engine.
   pip install fastapi-async-sqlalchemy
 ```
 
-### Important !!!
-If you use ```sqlmodel``` install ```sqlalchemy<=1.4.41```
+
+It also works with ```sqlmodel```
 
 
 ### Examples
@@ -157,9 +160,10 @@ app.add_middleware(
 routes.py
 
 ```python
+import asyncio
+
 from fastapi import APIRouter
-from sqlalchemy import column
-from sqlalchemy import table
+from sqlalchemy import column, table, text
 
 from databases import first_db, second_db
 
@@ -177,4 +181,22 @@ async def get_files_from_first_db():
 async def get_files_from_second_db():
     result = await second_db.session.execute(foo.select())
     return result.fetchall()
+
+
+@router.get("/concurrent-queries")
+async def parallel_select():
+    async with first_db(multi_sessions=True):
+        async def execute_query(query):
+            return await first_db.session.execute(text(query))
+
+        tasks = [
+            asyncio.create_task(execute_query("SELECT 1")),
+            asyncio.create_task(execute_query("SELECT 2")),
+            asyncio.create_task(execute_query("SELECT 3")),
+            asyncio.create_task(execute_query("SELECT 4")),
+            asyncio.create_task(execute_query("SELECT 5")),
+            asyncio.create_task(execute_query("SELECT 6")),
+        ]
+
+        await asyncio.gather(*tasks)
 ```

{fastapi-async-sqlalchemy-0.6.0 → fastapi_async_sqlalchemy-0.7.0.dev1}/README.md

@@ -5,7 +5,7 @@
 [![codecov](https://codecov.io/gh/h0rn3t/fastapi-async-sqlalchemy/branch/main/graph/badge.svg?token=F4NJ34WKPY)](https://codecov.io/gh/h0rn3t/fastapi-async-sqlalchemy)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 [![pip](https://img.shields.io/pypi/v/fastapi_async_sqlalchemy?color=blue)](https://pypi.org/project/fastapi-async-sqlalchemy/)
-[![Downloads](https://pepy.tech/badge/fastapi-async-sqlalchemy)](https://pepy.tech/project/fastapi-async-sqlalchemy)
+[![Downloads](https://static.pepy.tech/badge/fastapi-async-sqlalchemy)](https://pepy.tech/project/fastapi-async-sqlalchemy)
 [![Updates](https://pyup.io/repos/github/h0rn3t/fastapi-async-sqlalchemy/shield.svg)](https://pyup.io/repos/github/h0rn3t/fastapi-async-sqlalchemy/)
 
 ### Description
@@ -18,8 +18,8 @@ Provides SQLAlchemy middleware for FastAPI using AsyncSession and async engine.
   pip install fastapi-async-sqlalchemy
 ```
 
-### Important !!!
-If you use ```sqlmodel``` install ```sqlalchemy<=1.4.41```
+
+It also works with ```sqlmodel```
 
 
 ### Examples
@@ -127,9 +127,10 @@ app.add_middleware(
 routes.py
 
 ```python
+import asyncio
+
 from fastapi import APIRouter
-from sqlalchemy import column
-from sqlalchemy import table
+from sqlalchemy import column, table, text
 
 from databases import first_db, second_db
 
@@ -147,4 +148,22 @@ async def get_files_from_first_db():
 async def get_files_from_second_db():
     result = await second_db.session.execute(foo.select())
     return result.fetchall()
-```
+
+
+@router.get("/concurrent-queries")
+async def parallel_select():
+    async with first_db(multi_sessions=True):
+        async def execute_query(query):
+            return await first_db.session.execute(text(query))
+
+        tasks = [
+            asyncio.create_task(execute_query("SELECT 1")),
+            asyncio.create_task(execute_query("SELECT 2")),
+            asyncio.create_task(execute_query("SELECT 3")),
+            asyncio.create_task(execute_query("SELECT 4")),
+            asyncio.create_task(execute_query("SELECT 5")),
+            asyncio.create_task(execute_query("SELECT 6")),
+        ]
+
+        await asyncio.gather(*tasks)
+```

{fastapi-async-sqlalchemy-0.6.0 → fastapi_async_sqlalchemy-0.7.0.dev1}/fastapi_async_sqlalchemy/__init__.py

@@ -2,4 +2,4 @@ from fastapi_async_sqlalchemy.middleware import SQLAlchemyMiddleware, db
 
 __all__ = ["db", "SQLAlchemyMiddleware"]
 
-__version__ = "0.6.0"
+__version__ = "0.7.0.dev1"

fastapi_async_sqlalchemy-0.7.0.dev1/fastapi_async_sqlalchemy/middleware.py

@@ -0,0 +1,163 @@
+import asyncio
+from asyncio import Task
+from contextvars import ContextVar
+from typing import Dict, Optional, Union
+
+from sqlalchemy.engine import Engine
+from sqlalchemy.engine.url import URL
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import Request
+from starlette.types import ASGIApp
+
+from fastapi_async_sqlalchemy.exceptions import MissingSessionError, SessionNotInitialisedError
+
+try:
+    from sqlalchemy.ext.asyncio import async_sessionmaker  # noqa: F811
+except ImportError:
+    from sqlalchemy.orm import sessionmaker as async_sessionmaker
+
+
+def create_middleware_and_session_proxy():
+    _Session: Optional[async_sessionmaker] = None
+    _session: ContextVar[Optional[AsyncSession]] = ContextVar("_session", default=None)
+    _multi_sessions_ctx: ContextVar[bool] = ContextVar("_multi_sessions_context", default=False)
+    _commit_on_exit_ctx: ContextVar[bool] = ContextVar("_commit_on_exit_ctx", default=False)
+    # Usage of context vars inside closures is not recommended, since they are not properly
+    # garbage collected, but in our use case context var is created on program startup and
+    # is used throughout the whole its lifecycle.
+
+    class SQLAlchemyMiddleware(BaseHTTPMiddleware):
+        def __init__(
+            self,
+            app: ASGIApp,
+            db_url: Optional[Union[str, URL]] = None,
+            custom_engine: Optional[Engine] = None,
+            engine_args: Dict = None,
+            session_args: Dict = None,
+            commit_on_exit: bool = False,
+        ):
+            super().__init__(app)
+            self.commit_on_exit = commit_on_exit
+            engine_args = engine_args or {}
+            session_args = session_args or {}
+
+            if not custom_engine and not db_url:
+                raise ValueError("You need to pass a db_url or a custom_engine parameter.")
+            if not custom_engine:
+                engine = create_async_engine(db_url, **engine_args)
+            else:
+                engine = custom_engine
+
+            nonlocal _Session
+            _Session = async_sessionmaker(
+                engine, class_=AsyncSession, expire_on_commit=False, **session_args
+            )
+
+        async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
+            async with DBSession(commit_on_exit=self.commit_on_exit):
+                return await call_next(request)
+
+    class DBSessionMeta(type):
+        @property
+        def session(self) -> AsyncSession:
+            """Return an instance of Session local to the current async context."""
+            if _Session is None:
+                raise SessionNotInitialisedError
+
+            multi_sessions = _multi_sessions_ctx.get()
+            if multi_sessions:
+                """In this case, we need to create a new session for each task.
+                We also need to commit the session on exit if commit_on_exit is True.
+                This is useful when we need to run multiple queries in parallel.
+                For example, when we need to run multiple queries in parallel in a route handler.
+                Example:
+                ```python
+                    async with db(multi_sessions=True):
+                        async def execute_query(query):
+                            return await db.session.execute(text(query))
+
+                        tasks = [
+                            asyncio.create_task(execute_query("SELECT 1")),
+                            asyncio.create_task(execute_query("SELECT 2")),
+                            asyncio.create_task(execute_query("SELECT 3")),
+                            asyncio.create_task(execute_query("SELECT 4")),
+                            asyncio.create_task(execute_query("SELECT 5")),
+                            asyncio.create_task(execute_query("SELECT 6")),
+                        ]
+
+                        await asyncio.gather(*tasks)
+                ```
+                """
+                commit_on_exit = _commit_on_exit_ctx.get()
+                task: Task = asyncio.current_task()  # type: ignore
+                if not hasattr(task, "_db_session"):
+                    task._db_session = _Session()  # type: ignore
+
+                    def cleanup(future):
+                        session = getattr(task, "_db_session", None)
+                        if session:
+
+                            async def do_cleanup():
+                                try:
+                                    if future.exception():
+                                        await session.rollback()
+                                    else:
+                                        if commit_on_exit:
+                                            await session.commit()
+                                finally:
+                                    await session.close()
+
+                            asyncio.create_task(do_cleanup())
+
+                    task.add_done_callback(cleanup)
+                return task._db_session  # type: ignore
+            else:
+                session = _session.get()
+                if session is None:
+                    raise MissingSessionError
+                return session
+
+    class DBSession(metaclass=DBSessionMeta):
+        def __init__(
+            self,
+            session_args: Dict = None,
+            commit_on_exit: bool = False,
+            multi_sessions: bool = False,
+        ):
+            self.token = None
+            self.multi_sessions_token = None
+            self.commit_on_exit_token = None
+            self.session_args = session_args or {}
+            self.commit_on_exit = commit_on_exit
+            self.multi_sessions = multi_sessions
+
+        async def __aenter__(self):
+            if not isinstance(_Session, async_sessionmaker):
+                raise SessionNotInitialisedError
+
+            if self.multi_sessions:
+                self.multi_sessions_token = _multi_sessions_ctx.set(True)
+                self.commit_on_exit_token = _commit_on_exit_ctx.set(self.commit_on_exit)
+
+            self.token = _session.set(_Session(**self.session_args))
+            return type(self)
+
+        async def __aexit__(self, exc_type, exc_value, traceback):
+            session = _session.get()
+            try:
+                if exc_type is not None:
+                    await session.rollback()
+                elif self.commit_on_exit:
+                    await session.commit()
+            finally:
+                await session.close()
+                _session.reset(self.token)
+                if self.multi_sessions_token is not None:
+                    _multi_sessions_ctx.reset(self.multi_sessions_token)
+                    _commit_on_exit_ctx.reset(self.commit_on_exit_token)
+
+    return SQLAlchemyMiddleware, DBSession
+
+
+SQLAlchemyMiddleware, db = create_middleware_and_session_proxy()

{fastapi-async-sqlalchemy-0.6.0 → fastapi_async_sqlalchemy-0.7.0.dev1}/fastapi_async_sqlalchemy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastapi-async-sqlalchemy
-Version: 0.6.0
+Version: 0.7.0.dev1
 Summary: SQLAlchemy middleware for FastAPI
 Home-page: https://github.com/h0rn3t/fastapi-async-sqlalchemy.git
 Author: Eugene Shershen
@@ -8,7 +8,7 @@ Author-email: h0rn3t.null@gmail.com
 License: MIT
 Project-URL: Code, https://github.com/h0rn3t/fastapi-async-sqlalchemy
 Project-URL: Issue tracker, https://github.com/h0rn3t/fastapi-async-sqlalchemy/issues
-Classifier: Development Status :: 4 - Beta
+Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: Web Environment
 Classifier: Framework :: AsyncIO
 Classifier: Intended Audience :: Developers
@@ -19,6 +19,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Topic :: Internet :: WWW/HTTP :: HTTP Servers
@@ -27,6 +28,8 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.7
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: starlette>=0.13.6
+Requires-Dist: SQLAlchemy>=1.4.19
 
 # SQLAlchemy FastAPI middleware
 
@@ -35,7 +38,7 @@ License-File: LICENSE
 [![codecov](https://codecov.io/gh/h0rn3t/fastapi-async-sqlalchemy/branch/main/graph/badge.svg?token=F4NJ34WKPY)](https://codecov.io/gh/h0rn3t/fastapi-async-sqlalchemy)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 [![pip](https://img.shields.io/pypi/v/fastapi_async_sqlalchemy?color=blue)](https://pypi.org/project/fastapi-async-sqlalchemy/)
-[![Downloads](https://pepy.tech/badge/fastapi-async-sqlalchemy)](https://pepy.tech/project/fastapi-async-sqlalchemy)
+[![Downloads](https://static.pepy.tech/badge/fastapi-async-sqlalchemy)](https://pepy.tech/project/fastapi-async-sqlalchemy)
 [![Updates](https://pyup.io/repos/github/h0rn3t/fastapi-async-sqlalchemy/shield.svg)](https://pyup.io/repos/github/h0rn3t/fastapi-async-sqlalchemy/)
 
 ### Description
@@ -48,8 +51,8 @@ Provides SQLAlchemy middleware for FastAPI using AsyncSession and async engine.
   pip install fastapi-async-sqlalchemy
 ```
 
-### Important !!!
-If you use ```sqlmodel``` install ```sqlalchemy<=1.4.41```
+
+It also works with ```sqlmodel```
 
 
 ### Examples
@@ -157,9 +160,10 @@ app.add_middleware(
 routes.py
 
 ```python
+import asyncio
+
 from fastapi import APIRouter
-from sqlalchemy import column
-from sqlalchemy import table
+from sqlalchemy import column, table, text
 
 from databases import first_db, second_db
 
@@ -177,4 +181,22 @@ async def get_files_from_first_db():
 async def get_files_from_second_db():
     result = await second_db.session.execute(foo.select())
     return result.fetchall()
+
+
+@router.get("/concurrent-queries")
+async def parallel_select():
+    async with first_db(multi_sessions=True):
+        async def execute_query(query):
+            return await first_db.session.execute(text(query))
+
+        tasks = [
+            asyncio.create_task(execute_query("SELECT 1")),
+            asyncio.create_task(execute_query("SELECT 2")),
+            asyncio.create_task(execute_query("SELECT 3")),
+            asyncio.create_task(execute_query("SELECT 4")),
+            asyncio.create_task(execute_query("SELECT 5")),
+            asyncio.create_task(execute_query("SELECT 6")),
+        ]
+
+        await asyncio.gather(*tasks)
 ```

{fastapi-async-sqlalchemy-0.6.0 → fastapi_async_sqlalchemy-0.7.0.dev1}/fastapi_async_sqlalchemy.egg-info/SOURCES.txt

@@ -11,4 +11,5 @@ fastapi_async_sqlalchemy.egg-info/SOURCES.txt
 fastapi_async_sqlalchemy.egg-info/dependency_links.txt
 fastapi_async_sqlalchemy.egg-info/not-zip-safe
 fastapi_async_sqlalchemy.egg-info/requires.txt
-fastapi_async_sqlalchemy.egg-info/top_level.txt
+fastapi_async_sqlalchemy.egg-info/top_level.txt
+tests/test_session.py

{fastapi-async-sqlalchemy-0.6.0 → fastapi_async_sqlalchemy-0.7.0.dev1}/setup.py

@@ -29,7 +29,7 @@ setup(
     python_requires=">=3.7",
     install_requires=["starlette>=0.13.6", "SQLAlchemy>=1.4.19"],
     classifiers=[
-        "Development Status :: 4 - Beta",
+        "Development Status :: 5 - Production/Stable",
         "Environment :: Web Environment",
         "Framework :: AsyncIO",
         "Intended Audience :: Developers",
@@ -40,6 +40,7 @@ setup(
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
         "Programming Language :: Python :: 3 :: Only",
         "Programming Language :: Python :: Implementation :: CPython",
         "Topic :: Internet :: WWW/HTTP :: HTTP Servers",

fastapi_async_sqlalchemy-0.7.0.dev1/tests/test_session.py

@@ -0,0 +1,175 @@
+import asyncio
+
+import pytest
+from sqlalchemy import text
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from fastapi_async_sqlalchemy.exceptions import MissingSessionError, SessionNotInitialisedError
+
+db_url = "sqlite+aiosqlite://"
+
+
+@pytest.mark.asyncio
+async def test_init(app, SQLAlchemyMiddleware):
+    mw = SQLAlchemyMiddleware(app, db_url=db_url)
+    assert isinstance(mw, BaseHTTPMiddleware)
+
+
+@pytest.mark.asyncio
+async def test_init_required_args(app, SQLAlchemyMiddleware):
+    with pytest.raises(ValueError) as exc_info:
+        SQLAlchemyMiddleware(app)
+
+    assert exc_info.value.args[0] == "You need to pass a db_url or a custom_engine parameter."
+
+
+@pytest.mark.asyncio
+async def test_init_required_args_custom_engine(app, db, SQLAlchemyMiddleware):
+    custom_engine = create_async_engine(db_url)
+    SQLAlchemyMiddleware(app, custom_engine=custom_engine)
+
+
+@pytest.mark.asyncio
+async def test_init_correct_optional_args(app, db, SQLAlchemyMiddleware):
+    engine_args = {"echo": True}
+    # session_args = {"expire_on_commit": False}
+
+    SQLAlchemyMiddleware(app, db_url, engine_args=engine_args, session_args={})
+
+    async with db():
+        # assert not db.session.expire_on_commit
+        engine = db.session.bind
+        assert engine.echo
+
+    async with db() as db_ctx:
+        engine = db_ctx.session.bind
+        assert engine.echo
+
+
+@pytest.mark.asyncio
+async def test_init_incorrect_optional_args(app, SQLAlchemyMiddleware):
+    with pytest.raises(TypeError) as exc_info:
+        SQLAlchemyMiddleware(app, db_url=db_url, invalid_args="test")
+
+    assert "__init__() got an unexpected keyword argument 'invalid_args'" in exc_info.value.args[0]
+
+
+@pytest.mark.asyncio
+async def test_inside_route(app, client, db, SQLAlchemyMiddleware):
+    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url)
+
+    @app.get("/")
+    def test_get():
+        assert isinstance(db.session, AsyncSession)
+
+    client.get("/")
+
+
+@pytest.mark.asyncio
+async def test_inside_route_without_middleware_fails(app, client, db):
+    @app.get("/")
+    def test_get():
+        with pytest.raises(SessionNotInitialisedError):
+            db.session
+
+    client.get("/")
+
+
+@pytest.mark.asyncio
+async def test_outside_of_route(app, db, SQLAlchemyMiddleware):
+    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url)
+
+    async with db():
+        assert isinstance(db.session, AsyncSession)
+
+
+@pytest.mark.asyncio
+async def test_outside_of_route_without_middleware_fails(db):
+    with pytest.raises(SessionNotInitialisedError):
+        db.session
+
+    with pytest.raises(SessionNotInitialisedError):
+        async with db():
+            pass
+
+
+@pytest.mark.asyncio
+async def test_outside_of_route_without_context_fails(app, db, SQLAlchemyMiddleware):
+    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url)
+
+    with pytest.raises(MissingSessionError):
+        db.session
+
+
+@pytest.mark.asyncio
+async def test_init_session(app, db, SQLAlchemyMiddleware):
+    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url)
+
+    async with db():
+        assert isinstance(db.session, AsyncSession)
+
+
+@pytest.mark.asyncio
+async def test_db_session_commit_fail(app, db, SQLAlchemyMiddleware):
+    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url, commit_on_exit=True)
+
+    with pytest.raises(IntegrityError):
+        async with db():
+            raise IntegrityError("test", "test", "test")
+        db.session.close.assert_called_once()
+
+    async with db():
+        assert db.session
+
+
+@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):
+        async with db():
+            raise Exception
+
+        db.session.rollback.assert_called_once()
+
+
+@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)
+
+    session_args = {}
+
+    async with db(session_args=session_args, commit_on_exit=True):
+        assert isinstance(db.session, AsyncSession)
+
+    session_args = {"expire_on_commit": False}
+    async with db(session_args=session_args):
+        db.session
+
+
+@pytest.mark.asyncio
+async def test_multi_sessions(app, db, SQLAlchemyMiddleware):
+    app.add_middleware(SQLAlchemyMiddleware, db_url=db_url)
+
+    async with db(multi_sessions=True):
+
+        async def execute_query(query):
+            return await db.session.execute(text(query))
+
+        tasks = [
+            asyncio.create_task(execute_query("SELECT 1")),
+            asyncio.create_task(execute_query("SELECT 2")),
+            asyncio.create_task(execute_query("SELECT 3")),
+            asyncio.create_task(execute_query("SELECT 4")),
+            asyncio.create_task(execute_query("SELECT 5")),
+            asyncio.create_task(execute_query("SELECT 6")),
+        ]
+
+        res = await asyncio.gather(*tasks)
+        assert len(res) == 6

fastapi-async-sqlalchemy-0.6.0/fastapi_async_sqlalchemy/middleware.py

@@ -1,100 +0,0 @@
-from contextvars import ContextVar
-from typing import Dict, Optional, Union
-
-from sqlalchemy.engine import Engine
-from sqlalchemy.engine.url import URL
-from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
-from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
-from starlette.requests import Request
-from starlette.types import ASGIApp
-
-from fastapi_async_sqlalchemy.exceptions import MissingSessionError, SessionNotInitialisedError
-
-try:
-    from sqlalchemy.ext.asyncio import async_sessionmaker
-except ImportError:
-    from sqlalchemy.orm import sessionmaker as async_sessionmaker
-
-
-def create_middleware_and_session_proxy():
-    _Session: Optional[async_sessionmaker] = None
-    # Usage of context vars inside closures is not recommended, since they are not properly
-    # garbage collected, but in our use case context var is created on program startup and
-    # is used throughout the whole its lifecycle.
-    _session: ContextVar[Optional[AsyncSession]] = ContextVar("_session", default=None)
-
-    class SQLAlchemyMiddleware(BaseHTTPMiddleware):
-        def __init__(
-            self,
-            app: ASGIApp,
-            db_url: Optional[Union[str, URL]] = None,
-            custom_engine: Optional[Engine] = None,
-            engine_args: Dict = None,
-            session_args: Dict = None,
-            commit_on_exit: bool = False,
-        ):
-            super().__init__(app)
-            self.commit_on_exit = commit_on_exit
-            engine_args = engine_args or {}
-            session_args = session_args or {}
-
-            if not custom_engine and not db_url:
-                raise ValueError("You need to pass a db_url or a custom_engine parameter.")
-            if not custom_engine:
-                engine = create_async_engine(db_url, **engine_args)
-            else:
-                engine = custom_engine
-
-            nonlocal _Session
-            _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):
-        @property
-        def session(self) -> AsyncSession:
-            """Return an instance of Session local to the current async context."""
-            if _Session is None:
-                raise SessionNotInitialisedError
-
-            session = _session.get()
-            if session is None:
-                raise MissingSessionError
-
-            return session
-
-    class DBSession(metaclass=DBSessionMeta):
-        def __init__(self, session_args: Dict = None, commit_on_exit: bool = False):
-            self.token = None
-            self.session_args = session_args or {}
-            self.commit_on_exit = commit_on_exit
-
-        async def __aenter__(self):
-            if not isinstance(_Session, async_sessionmaker):
-                raise SessionNotInitialisedError
-
-            self.token = _session.set(_Session(**self.session_args))  # type: ignore
-            return type(self)
-
-        async def __aexit__(self, exc_type, exc_value, traceback):
-            session = _session.get()
-
-            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)
-
-    return SQLAlchemyMiddleware, DBSession
-
-
-SQLAlchemyMiddleware, db = create_middleware_and_session_proxy()