From d40c93e062a80989b7b9d8e2825dbac341b63fa3 Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 6 Sep 2023 13:50:29 +0300 Subject: [PATCH 01/34] Update README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 56d3a58..7629d15 100644 --- a/README.md +++ b/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 @@ -147,4 +147,4 @@ 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() -``` \ No newline at end of file +``` From 3f98f13743f8ff6429013565de845b4ac137de85 Mon Sep 17 00:00:00 2001 From: Vladislav Shepilov Date: Thu, 21 Dec 2023 21:28:40 +0300 Subject: [PATCH 02/34] bug fix for multiple databases support, fix #17 --- fastapi_async_sqlalchemy/middleware.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index 1a407fa..bb6024d 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -51,7 +51,7 @@ def __init__( ) async def dispatch(self, request: Request, call_next: RequestResponseEndpoint): - async with db(commit_on_exit=self.commit_on_exit): + async with DBSession(commit_on_exit=self.commit_on_exit): return await call_next(request) class DBSessionMeta(type): From cb52c9dfb25910aa745ef2762bf3223de943919a Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Wed, 17 Jan 2024 10:04:15 +0200 Subject: [PATCH 03/34] pytest_asyncio --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index bb00aee..4392846 100644 --- a/requirements.txt +++ b/requirements.txt @@ -38,5 +38,5 @@ urllib3>=1.25.9 wcwidth==0.1.7 zipp==3.1.0 black==22.3.0 -pytest-asyncio>=0.15.0 +pytest-asyncio==0.21.0 greenlet==2.0.2 From 992f9b8ed1a2acd144575128bd3d2a148aabddee Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Wed, 17 Jan 2024 10:05:32 +0200 Subject: [PATCH 04/34] pydantic==1.10.13 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 4392846..7c6581a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,7 +17,7 @@ packaging==19.2 pathspec>=0.9.0 pluggy==0.13.0 pycodestyle==2.5.0 -pydantic==1.10.2 +pydantic==1.10.13 pyflakes==2.1.1 pyparsing==2.4.2 pytest==7.2.0 From cb98d48dcf2ed206e6027ed3f0b75e52d86015a5 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Wed, 17 Jan 2024 10:08:46 +0200 Subject: [PATCH 05/34] pytest fixes --- requirements.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/requirements.txt b/requirements.txt index 7c6581a..0b34545 100644 --- a/requirements.txt +++ b/requirements.txt @@ -20,8 +20,8 @@ pycodestyle==2.5.0 pydantic==1.10.13 pyflakes==2.1.1 pyparsing==2.4.2 -pytest==7.2.0 -pytest-cov==2.11.1 +pytest>=7.2.0 +pytest-cov>=2.11.1 PyYAML>=5.4 regex==2020.2.20 requests>=2.22.0 @@ -38,5 +38,5 @@ urllib3>=1.25.9 wcwidth==0.1.7 zipp==3.1.0 black==22.3.0 -pytest-asyncio==0.21.0 +pytest-asyncio>=0.21.0 greenlet==2.0.2 From e3b9f834814d565e3c3d3f04dd4edbe208932147 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Wed, 17 Jan 2024 10:11:19 +0200 Subject: [PATCH 06/34] pytest fixes --- requirements.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/requirements.txt b/requirements.txt index 0b34545..7c6581a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -20,8 +20,8 @@ pycodestyle==2.5.0 pydantic==1.10.13 pyflakes==2.1.1 pyparsing==2.4.2 -pytest>=7.2.0 -pytest-cov>=2.11.1 +pytest==7.2.0 +pytest-cov==2.11.1 PyYAML>=5.4 regex==2020.2.20 requests>=2.22.0 @@ -38,5 +38,5 @@ urllib3>=1.25.9 wcwidth==0.1.7 zipp==3.1.0 black==22.3.0 -pytest-asyncio>=0.21.0 +pytest-asyncio==0.21.0 greenlet==2.0.2 From 0bd9197959ee9cfecaa318d8bbf7e00176c7c2ae Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 17 Jan 2024 10:19:15 +0200 Subject: [PATCH 07/34] Update __init__.py bump 0.6.1 --- fastapi_async_sqlalchemy/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastapi_async_sqlalchemy/__init__.py b/fastapi_async_sqlalchemy/__init__.py index 0a44930..ab2c3b7 100644 --- a/fastapi_async_sqlalchemy/__init__.py +++ b/fastapi_async_sqlalchemy/__init__.py @@ -2,4 +2,4 @@ __all__ = ["db", "SQLAlchemyMiddleware"] -__version__ = "0.6.0" +__version__ = "0.6.1" From 3239a5441dc75cef279bc55cb04f1913731a1db4 Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 17 Jan 2024 10:40:41 +0200 Subject: [PATCH 08/34] Update python-publish.yml --- .github/workflows/python-publish.yml | 32 +++++++++++++++++----------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 1a03a7b..bdaab28 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -1,11 +1,19 @@ # This workflow will upload a Python Package using Twine when a release is created -# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries +# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries + +# This workflow uses actions that are not certified by GitHub. +# They are provided by a third-party and are governed by +# separate terms of service, privacy policy, and support +# documentation. name: Upload Python Package on: release: - types: [created] + types: [published] + +permissions: + contents: read jobs: deploy: @@ -13,19 +21,19 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Set up Python - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: '3.x' - name: Install dependencies run: | python -m pip install --upgrade pip - pip install setuptools wheel twine - - name: Build and publish - env: - TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} - TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} - run: | - python setup.py sdist bdist_wheel - twine upload dist/* + pip install build + - name: Build package + run: python -m build + - name: Publish package + uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 + with: + user: __token__ + password: ${{ secrets.PYPI_API_TOKEN }} From 9c27807416f69aacd052a5722353548e6373c250 Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 24 Jan 2024 09:44:38 +0200 Subject: [PATCH 09/34] Update pyproject.toml --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 5a9141f..c48afa5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,3 +17,6 @@ include_trailing_comma = true force_grid_wrap = 0 use_parentheses = true line_length = 100 + +[tool.coverage.run] +concurrency = ["greenlet"] From 843cf93e36b17f5987adc1d9e4203c73cd9b741a Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 24 Jan 2024 09:47:07 +0200 Subject: [PATCH 10/34] Update pyproject.toml --- pyproject.toml | 3 --- 1 file changed, 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c48afa5..5a9141f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,3 @@ include_trailing_comma = true force_grid_wrap = 0 use_parentheses = true line_length = 100 - -[tool.coverage.run] -concurrency = ["greenlet"] From e3d49dd062b8c2c70523ea556606f449ef39c5fb Mon Sep 17 00:00:00 2001 From: Eugene Date: Mon, 12 Feb 2024 12:08:37 +0200 Subject: [PATCH 11/34] Update README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 7629d15..84546be 100644 --- a/README.md +++ b/README.md @@ -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 From 70704c3dfa2d25640e4b26136c00b170a87cb8a9 Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 21 Feb 2024 13:18:35 +0200 Subject: [PATCH 12/34] Update setup.py --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index d9b0d89..efb6712 100644 --- a/setup.py +++ b/setup.py @@ -40,6 +40,7 @@ "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", From 15a798b0bb13d2e13726c13f3bb9c0ebd5a5a182 Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 21 Feb 2024 13:19:45 +0200 Subject: [PATCH 13/34] Update setup.py --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index efb6712..09786b6 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ 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", From 4e7a673ed5fa1a552a1c9d157ce50166b085ba48 Mon Sep 17 00:00:00 2001 From: Eugene Date: Thu, 21 Mar 2024 00:05:42 +0200 Subject: [PATCH 14/34] Update requirements.txt --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 7c6581a..5bafbaa 100644 --- a/requirements.txt +++ b/requirements.txt @@ -29,7 +29,7 @@ httpx>=0.20.0 six==1.12.0 SQLAlchemy>=1.4.19 asyncpg>=0.27.0 -aiosqlite==0.19.0 +aiosqlite==0.20.0 sqlparse==0.4.4 starlette>=0.13.6 toml>=0.10.1 From 8324bdad0cce5733654a0a29fad4b5f0124defd0 Mon Sep 17 00:00:00 2001 From: Eugene Date: Fri, 26 Jul 2024 01:06:55 +0300 Subject: [PATCH 15/34] Update requirements.txt --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 5bafbaa..c6f0d4b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -30,7 +30,7 @@ six==1.12.0 SQLAlchemy>=1.4.19 asyncpg>=0.27.0 aiosqlite==0.20.0 -sqlparse==0.4.4 +sqlparse==0.5.1 starlette>=0.13.6 toml>=0.10.1 typed-ast>=1.4.2 From f247a7f25fb23e75ae003b328ec6d7dbc3aa2d42 Mon Sep 17 00:00:00 2001 From: Eugene Date: Fri, 26 Jul 2024 01:08:38 +0300 Subject: [PATCH 16/34] Update requirements.txt --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index c6f0d4b..d103a71 100644 --- a/requirements.txt +++ b/requirements.txt @@ -37,6 +37,6 @@ typed-ast>=1.4.2 urllib3>=1.25.9 wcwidth==0.1.7 zipp==3.1.0 -black==22.3.0 +black==24.4.2 pytest-asyncio==0.21.0 greenlet==2.0.2 From fd8bf11342176608672ee21a2b05e3dd47101013 Mon Sep 17 00:00:00 2001 From: Eugene Date: Fri, 26 Jul 2024 20:01:08 +0300 Subject: [PATCH 17/34] Update requirements.txt --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index d103a71..0279928 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,7 +13,7 @@ importlib-metadata==1.5.0 isort==4.3.21 mccabe==0.6.1 more-itertools==7.2.0 -packaging==19.2 +packaging>=22.0 pathspec>=0.9.0 pluggy==0.13.0 pycodestyle==2.5.0 From 8639280cf454c2800abeb456586fd0164af9f55a Mon Sep 17 00:00:00 2001 From: Eugene Date: Tue, 6 Aug 2024 13:08:51 +0300 Subject: [PATCH 18/34] Update ci.yml --- .github/workflows/ci.yml | 27 ++++++++++++--------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bd1026a..71b005a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,20 +13,17 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - build: [linux_3.8, windows_3.8, mac_3.8, linux_3.7] + build: [linux_3.9, windows_3.9, mac_3.9] include: - - build: linux_3.8 + - build: linux_3.9 os: ubuntu-latest - python: 3.8 - - build: windows_3.8 + python: 3.9 + - build: windows_3.9 os: windows-latest - python: 3.8 - - build: mac_3.8 + python: 3.9 + - build: mac_3.9 os: macos-latest - python: 3.8 - - build: linux_3.7 - os: ubuntu-latest - python: 3.7 + python: 3.9 steps: - name: Checkout repository uses: actions/checkout@v2 @@ -43,16 +40,16 @@ jobs: # test all the builds apart from linux_3.8... - name: Test with pytest - if: matrix.build != 'linux_3.8' + if: matrix.build != 'linux_3.9' run: pytest # only do the test coverage for linux_3.8 - name: Produce coverage report - if: matrix.build == 'linux_3.8' + if: matrix.build == 'linux_3.9' run: pytest --cov=fastapi_async_sqlalchemy --cov-report=xml - name: Upload coverage report - if: matrix.build == 'linux_3.8' + if: matrix.build == 'linux_3.9' uses: codecov/codecov-action@v1 with: file: ./coverage.xml @@ -67,7 +64,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v1 with: - python-version: 3.7 + python-version: 3.9 - name: Install dependencies run: pip install flake8 @@ -85,7 +82,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v1 with: - python-version: 3.7 + python-version: 3.9 - name: Install dependencies # isort needs all of the packages to be installed so it can From 698533958e028bc5fb6befe04470058ee850b625 Mon Sep 17 00:00:00 2001 From: Eugene Date: Tue, 6 Aug 2024 13:10:52 +0300 Subject: [PATCH 19/34] Update ci.yml --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 71b005a..b4395f2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -62,7 +62,7 @@ jobs: uses: actions/checkout@v2 - name: Set up Python - uses: actions/setup-python@v1 + uses: actions/setup-python@v4 with: python-version: 3.9 @@ -80,7 +80,7 @@ jobs: uses: actions/checkout@v2 - name: Set up Python - uses: actions/setup-python@v1 + uses: actions/setup-python@v4 with: python-version: 3.9 From 52272ca7eed67483901782af7688d5c5c6b1edea Mon Sep 17 00:00:00 2001 From: Eugene Date: Tue, 6 Aug 2024 13:12:40 +0300 Subject: [PATCH 20/34] Update ci.yml --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b4395f2..808c1c7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -29,7 +29,7 @@ jobs: uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python }} - uses: actions/setup-python@v1 + uses: actions/setup-python@v4 with: python-version: ${{ matrix.python }} From e6dc0a7ca86fa17f9826b22b12337ac33c1892ff Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 21:20:29 +0300 Subject: [PATCH 21/34] task-local sessions --- .pre-commit-config.yaml | 2 +- fastapi_async_sqlalchemy/middleware.py | 108 +++++++++++++++++++------ requirements.txt | 6 +- 3 files changed, 89 insertions(+), 27 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7723dcc..a21adc0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -41,7 +41,7 @@ repos: - --max-line-length=100 - --ignore=E203, E501, W503 - repo: https://github.com/pre-commit/mirrors-mypy - rev: v0.982 + rev: v1.12.0 hooks: - id: mypy additional_dependencies: diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index bb6024d..e059b7c 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -1,27 +1,21 @@ +import asyncio 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 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 -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) + _multi_sessions_ctx: ContextVar[bool] = ContextVar("_multi_sessions_context", default=False) class SQLAlchemyMiddleware(BaseHTTPMiddleware): def __init__( @@ -29,8 +23,8 @@ def __init__( app: ASGIApp, db_url: Optional[Union[str, URL]] = None, custom_engine: Optional[Engine] = None, - engine_args: Dict = None, - session_args: Dict = None, + engine_args: Optional[Dict] = None, + session_args: Optional[Dict] = None, commit_on_exit: bool = False, ): super().__init__(app) @@ -57,28 +51,96 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint): 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: + current_session = _session.get() + if current_session is None: raise MissingSessionError - return session + multi_sessions = _multi_sessions_ctx.get() + if multi_sessions: + # Если multi_sessions=True, используем Task-локальные сессии + task = asyncio.current_task() + if not hasattr(task, "_db_session"): + task._db_session = _Session() + + def cleanup(future): + session = getattr(task, "_db_session", None) + if session: + + async def do_cleanup(): + try: + if future.exception(): + await session.rollback() + else: + await session.commit() + finally: + await session.close() + + asyncio.create_task(do_cleanup()) + + task.add_done_callback(cleanup) + return task._db_session + else: + return current_session + + def __call__( + cls, + session_args: Optional[Dict] = None, + commit_on_exit: bool = False, + multi_sessions: bool = False, + ): + return cls._context_manager( + session_args=session_args, + commit_on_exit=commit_on_exit, + multi_sessions=multi_sessions, + ) + + def _context_manager( + cls, + session_args: Dict = None, + commit_on_exit: bool = False, + multi_sessions: bool = False, + ): + return DBSessionContextManager( + session_args=session_args, + commit_on_exit=commit_on_exit, + multi_sessions=multi_sessions, + ) class DBSession(metaclass=DBSessionMeta): - def __init__(self, session_args: Dict = None, commit_on_exit: bool = False): - self.token = None + pass + + class DBSessionContextManager: + def __init__( + self, + session_args: Dict = None, + commit_on_exit: bool = False, + multi_sessions: bool = False, + ): self.session_args = session_args or {} self.commit_on_exit = commit_on_exit + self.multi_sessions = multi_sessions + self.token = None + self.multi_sessions_token = None + self._session = None async def __aenter__(self): - if not isinstance(_Session, async_sessionmaker): + if _Session is None: raise SessionNotInitialisedError - self.token = _session.set(_Session(**self.session_args)) # type: ignore - return type(self) + if self.multi_sessions: + self.multi_sessions_token = _multi_sessions_ctx.set(True) + + self._session = _Session(**self.session_args) + self.token = _session.set(self._session) + return self + + @property + def session(self): + return self._session async def __aexit__(self, exc_type, exc_value, traceback): session = _session.get() @@ -86,13 +148,13 @@ async def __aexit__(self, exc_type, exc_value, traceback): 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 + 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) return SQLAlchemyMiddleware, DBSession diff --git a/requirements.txt b/requirements.txt index 0279928..d3232d8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,13 +17,13 @@ packaging>=22.0 pathspec>=0.9.0 pluggy==0.13.0 pycodestyle==2.5.0 -pydantic==1.10.13 +pydantic==1.10.18 pyflakes==2.1.1 pyparsing==2.4.2 pytest==7.2.0 pytest-cov==2.11.1 PyYAML>=5.4 -regex==2020.2.20 +regex>=2020.2.20 requests>=2.22.0 httpx>=0.20.0 six==1.12.0 @@ -39,4 +39,4 @@ wcwidth==0.1.7 zipp==3.1.0 black==24.4.2 pytest-asyncio==0.21.0 -greenlet==2.0.2 +greenlet==3.1.1 From 5af8f493f336e08672ac23f766b3928140c8261f Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 21:20:59 +0300 Subject: [PATCH 22/34] task-local sessions --- fastapi_async_sqlalchemy/middleware.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index e059b7c..229c5fa 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -51,7 +51,6 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint): class DBSessionMeta(type): @property def session(self) -> AsyncSession: - """Возвращает текущую сессию из контекста.""" if _Session is None: raise SessionNotInitialisedError @@ -61,7 +60,6 @@ def session(self) -> AsyncSession: multi_sessions = _multi_sessions_ctx.get() if multi_sessions: - # Если multi_sessions=True, используем Task-локальные сессии task = asyncio.current_task() if not hasattr(task, "_db_session"): task._db_session = _Session() From 7aa6decca560dd0abe41557b44aec8319212c7e0 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 21:28:16 +0300 Subject: [PATCH 23/34] commit_on_exit fix --- fastapi_async_sqlalchemy/middleware.py | 70 +++++++++----------------- 1 file changed, 23 insertions(+), 47 deletions(-) diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index 229c5fa..195d0ac 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -11,11 +11,17 @@ 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 _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) class SQLAlchemyMiddleware(BaseHTTPMiddleware): def __init__( @@ -23,8 +29,8 @@ def __init__( app: ASGIApp, db_url: Optional[Union[str, URL]] = None, custom_engine: Optional[Engine] = None, - engine_args: Optional[Dict] = None, - session_args: Optional[Dict] = None, + engine_args: Dict = None, + session_args: Dict = None, commit_on_exit: bool = False, ): super().__init__(app) @@ -54,12 +60,9 @@ def session(self) -> AsyncSession: if _Session is None: raise SessionNotInitialisedError - current_session = _session.get() - if current_session is None: - raise MissingSessionError - multi_sessions = _multi_sessions_ctx.get() if multi_sessions: + commit_on_exit = _commit_on_exit_ctx.get() task = asyncio.current_task() if not hasattr(task, "_db_session"): task._db_session = _Session() @@ -73,7 +76,8 @@ async def do_cleanup(): if future.exception(): await session.rollback() else: - await session.commit() + if commit_on_exit: + await session.commit() finally: await session.close() @@ -82,67 +86,38 @@ async def do_cleanup(): task.add_done_callback(cleanup) return task._db_session else: - return current_session - - def __call__( - cls, - session_args: Optional[Dict] = None, - commit_on_exit: bool = False, - multi_sessions: bool = False, - ): - return cls._context_manager( - session_args=session_args, - commit_on_exit=commit_on_exit, - multi_sessions=multi_sessions, - ) - - def _context_manager( - cls, - session_args: Dict = None, - commit_on_exit: bool = False, - multi_sessions: bool = False, - ): - return DBSessionContextManager( - session_args=session_args, - commit_on_exit=commit_on_exit, - multi_sessions=multi_sessions, - ) + session = _session.get() + if session is None: + raise MissingSessionError + return session class DBSession(metaclass=DBSessionMeta): - pass - - class DBSessionContextManager: 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 - self.token = None - self.multi_sessions_token = None - self._session = None async def __aenter__(self): - if _Session is None: + 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._session = _Session(**self.session_args) - self.token = _session.set(self._session) - return self - - @property - def session(self): - return self._session + 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() @@ -153,6 +128,7 @@ async def __aexit__(self, exc_type, exc_value, traceback): _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 From d95a2b087bd904d52697dbace3a4a31b3209cf49 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 22:08:20 +0300 Subject: [PATCH 24/34] added test_multi_sessions --- tests/test_session.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/test_session.py b/tests/test_session.py index 06163ac..c7e8666 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -1,4 +1,7 @@ +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 @@ -148,3 +151,24 @@ async def test_db_context_session_args(app, db, SQLAlchemyMiddleware, commit_on_ 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 From cb29f3e67f26e0d4d2a8c2b56df0aadf59866bb4 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 22:14:11 +0300 Subject: [PATCH 25/34] added test_multi_sessions --- README.md | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 84546be..70f1c6a 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ +from pytest_cov.embed import multiprocessing_finishfrom tests.test_session import db_url + # SQLAlchemy FastAPI middleware [![ci](https://img.shields.io/badge/Support-Ukraine-FFD500?style=flat&labelColor=005BBB)](https://img.shields.io/badge/Support-Ukraine-FFD500?style=flat&labelColor=005BBB) @@ -127,9 +129,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 +150,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) ``` From cb081f1e7b98f5758a13031116cb56d512582946 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 22:27:48 +0300 Subject: [PATCH 26/34] fixes mypy --- .pre-commit-config.yaml | 2 +- fastapi_async_sqlalchemy/middleware.py | 9 +++++---- tests/test_session.py | 1 + 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a21adc0..7723dcc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -41,7 +41,7 @@ repos: - --max-line-length=100 - --ignore=E203, E501, W503 - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.12.0 + rev: v0.982 hooks: - id: mypy additional_dependencies: diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index 195d0ac..6893e28 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -1,4 +1,5 @@ import asyncio +from asyncio import Task from contextvars import ContextVar from typing import Dict, Optional, Union @@ -12,7 +13,7 @@ from fastapi_async_sqlalchemy.exceptions import MissingSessionError, SessionNotInitialisedError try: - from sqlalchemy.ext.asyncio import async_sessionmaker + from sqlalchemy.ext.asyncio import async_sessionmaker # noqa: F811 except ImportError: from sqlalchemy.orm import sessionmaker as async_sessionmaker @@ -63,9 +64,9 @@ def session(self) -> AsyncSession: multi_sessions = _multi_sessions_ctx.get() if multi_sessions: commit_on_exit = _commit_on_exit_ctx.get() - task = asyncio.current_task() + task: Task = asyncio.current_task() # type: ignore if not hasattr(task, "_db_session"): - task._db_session = _Session() + task._db_session = _Session() # type: ignore def cleanup(future): session = getattr(task, "_db_session", None) @@ -84,7 +85,7 @@ async def do_cleanup(): asyncio.create_task(do_cleanup()) task.add_done_callback(cleanup) - return task._db_session + return task._db_session # type: ignore else: session = _session.get() if session is None: diff --git a/tests/test_session.py b/tests/test_session.py index c7e8666..82f5dc9 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -158,6 +158,7 @@ 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)) From 5c9f01fa0c2ae247c112ca98dc7a24fd350bccbe Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 22:28:41 +0300 Subject: [PATCH 27/34] fixes mypy --- README.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/README.md b/README.md index 70f1c6a..6d1e722 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,3 @@ -from pytest_cov.embed import multiprocessing_finishfrom tests.test_session import db_url - # SQLAlchemy FastAPI middleware [![ci](https://img.shields.io/badge/Support-Ukraine-FFD500?style=flat&labelColor=005BBB)](https://img.shields.io/badge/Support-Ukraine-FFD500?style=flat&labelColor=005BBB) From aafe26b517e4a3e53428574b182ac8f27df276e6 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 22:39:11 +0300 Subject: [PATCH 28/34] fixes mypy --- fastapi_async_sqlalchemy/middleware.py | 27 ++++++++++++++++++++++++++ requirements.txt | 4 ++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index 6893e28..f70a49e 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -23,6 +23,9 @@ def create_middleware_and_session_proxy(): _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__( @@ -58,11 +61,35 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint): 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: + """If multi_sessions is True, we are in a context where multiple sessions are allowed. + 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"): diff --git a/requirements.txt b/requirements.txt index d3232d8..e3a0644 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,7 +8,7 @@ coverage>=5.2.1 entrypoints==0.3 fastapi==0.90.0 # pyup: ignore flake8==3.7.9 -idna==2.8 +idna==3.7 importlib-metadata==1.5.0 isort==4.3.21 mccabe==0.6.1 @@ -36,7 +36,7 @@ toml>=0.10.1 typed-ast>=1.4.2 urllib3>=1.25.9 wcwidth==0.1.7 -zipp==3.1.0 +zipp==3.19.1 black==24.4.2 pytest-asyncio==0.21.0 greenlet==3.1.1 From a657d8f2067f6e1b999bd8522ed1731ab9122505 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 22:57:27 +0300 Subject: [PATCH 29/34] fixes mypy --- fastapi_async_sqlalchemy/middleware.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index f70a49e..06f3268 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -67,8 +67,7 @@ def session(self) -> AsyncSession: multi_sessions = _multi_sessions_ctx.get() if multi_sessions: - """If multi_sessions is True, we are in a context where multiple sessions are allowed. - In this case, we need to create a new session for each task. + """ 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. From 65ed17a3ce8f758b728d3e1cb00e7747815671cb Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 22:58:51 +0300 Subject: [PATCH 30/34] fixes mypy --- fastapi_async_sqlalchemy/middleware.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index 06f3268..e3b1563 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -67,7 +67,7 @@ def session(self) -> AsyncSession: multi_sessions = _multi_sessions_ctx.get() if multi_sessions: - """ In this case, we need to create a new session for each task. + """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. From 83d7bef8b99c83493444aa24bfc8f64094186220 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Thu, 17 Oct 2024 23:45:44 +0300 Subject: [PATCH 31/34] fixes mypy --- fastapi_async_sqlalchemy/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastapi_async_sqlalchemy/__init__.py b/fastapi_async_sqlalchemy/__init__.py index ab2c3b7..21c6edc 100644 --- a/fastapi_async_sqlalchemy/__init__.py +++ b/fastapi_async_sqlalchemy/__init__.py @@ -2,4 +2,4 @@ __all__ = ["db", "SQLAlchemyMiddleware"] -__version__ = "0.6.1" +__version__ = "0.7.0.dev1" From 0c74aaf2fc9808374e5cc148a96a32f38f3315ac Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Fri, 18 Oct 2024 13:04:31 +0300 Subject: [PATCH 32/34] WIP: multi_sessions --- fastapi_async_sqlalchemy/__init__.py | 2 +- fastapi_async_sqlalchemy/middleware.py | 75 +++++++++++++------------- tests/test_session.py | 25 +++++++++ 3 files changed, 64 insertions(+), 38 deletions(-) diff --git a/fastapi_async_sqlalchemy/__init__.py b/fastapi_async_sqlalchemy/__init__.py index 21c6edc..2821653 100644 --- a/fastapi_async_sqlalchemy/__init__.py +++ b/fastapi_async_sqlalchemy/__init__.py @@ -2,4 +2,4 @@ __all__ = ["db", "SQLAlchemyMiddleware"] -__version__ = "0.7.0.dev1" +__version__ = "0.7.0.dev2" diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index e3b1563..13cb359 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -1,5 +1,4 @@ import asyncio -from asyncio import Task from contextvars import ContextVar from typing import Dict, Optional, Union @@ -22,6 +21,9 @@ 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) + _task_session_ctx: ContextVar[Optional[AsyncSession]] = ContextVar( + "_task_session_ctx", default=None + ) _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 @@ -90,28 +92,26 @@ async def execute_query(query): ``` """ 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 + session = _task_session_ctx.get() + if session is None: + session = _Session() + _task_session_ctx.set(session) + + async def cleanup(): + try: + if commit_on_exit: + await session.commit() + except Exception: + await session.rollback() + raise + finally: + await session.close() + _task_session_ctx.set(None) + + task = asyncio.current_task() + if task is not None: + task.add_done_callback(lambda t: asyncio.create_task(cleanup())) + return session else: session = _session.get() if session is None: @@ -139,23 +139,24 @@ async def __aenter__(self): 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)) + else: + 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) + if self.multi_sessions: + _multi_sessions_ctx.reset(self.multi_sessions_token) + _commit_on_exit_ctx.reset(self.commit_on_exit_token) + else: + 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) return SQLAlchemyMiddleware, DBSession diff --git a/tests/test_session.py b/tests/test_session.py index 82f5dc9..9400fea 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -173,3 +173,28 @@ async def execute_query(query): res = await asyncio.gather(*tasks) assert len(res) == 6 + + +@pytest.mark.asyncio +async def test_concurrent_inserts(app, db, SQLAlchemyMiddleware): + app.add_middleware(SQLAlchemyMiddleware, db_url=db_url) + + async with db(multi_sessions=True, commit_on_exit=True): + await db.session.execute( + text("CREATE TABLE IF NOT EXISTS my_model (id INTEGER PRIMARY KEY, value TEXT)") + ) + + async def insert_data(value): + await db.session.execute( + text("INSERT INTO my_model (value) VALUES (:value)"), {"value": value} + ) + await db.session.flush() + + tasks = [asyncio.create_task(insert_data(f"value_{i}")) for i in range(10)] + + result_ids = await asyncio.gather(*tasks) + assert len(result_ids) == 10 + + records = await db.session.execute(text("SELECT * FROM my_model")) + records = records.scalars().all() + assert len(records) == 10 From 955f538fdd908f3fdae601ac52d798c6eac1e937 Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Sat, 19 Oct 2024 01:12:50 +0300 Subject: [PATCH 33/34] WIP: multi_sessions --- fastapi_async_sqlalchemy/__init__.py | 2 +- fastapi_async_sqlalchemy/middleware.py | 39 +++++++++++--------------- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/fastapi_async_sqlalchemy/__init__.py b/fastapi_async_sqlalchemy/__init__.py index 2821653..8afd39f 100644 --- a/fastapi_async_sqlalchemy/__init__.py +++ b/fastapi_async_sqlalchemy/__init__.py @@ -2,4 +2,4 @@ __all__ = ["db", "SQLAlchemyMiddleware"] -__version__ = "0.7.0.dev2" +__version__ = "0.7.0.dev3" diff --git a/fastapi_async_sqlalchemy/middleware.py b/fastapi_async_sqlalchemy/middleware.py index 13cb359..1171ede 100644 --- a/fastapi_async_sqlalchemy/middleware.py +++ b/fastapi_async_sqlalchemy/middleware.py @@ -21,9 +21,6 @@ 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) - _task_session_ctx: ContextVar[Optional[AsyncSession]] = ContextVar( - "_task_session_ctx", default=None - ) _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 @@ -92,25 +89,22 @@ async def execute_query(query): ``` """ commit_on_exit = _commit_on_exit_ctx.get() - session = _task_session_ctx.get() - if session is None: - session = _Session() - _task_session_ctx.set(session) - - async def cleanup(): - try: - if commit_on_exit: - await session.commit() - except Exception: - await session.rollback() - raise - finally: - await session.close() - _task_session_ctx.set(None) - - task = asyncio.current_task() - if task is not None: - task.add_done_callback(lambda t: asyncio.create_task(cleanup())) + # Always create a new session for each access when multi_sessions=True + session = _Session() + + async def cleanup(): + try: + if commit_on_exit: + await session.commit() + except Exception: + await session.rollback() + raise + finally: + await session.close() + + task = asyncio.current_task() + if task is not None: + task.add_done_callback(lambda t: asyncio.create_task(cleanup())) return session else: session = _session.get() @@ -126,7 +120,6 @@ def __init__( 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 From a76b7bb67ea6e1fe90f342bbbe002d58ec1f876f Mon Sep 17 00:00:00 2001 From: Eugene Shershen Date: Fri, 17 Jan 2025 10:47:00 +0200 Subject: [PATCH 34/34] fix import create_middleware_and_session_proxy --- README.md | 2 +- fastapi_async_sqlalchemy/__init__.py | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 6d1e722..32898a2 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ Provides SQLAlchemy middleware for FastAPI using AsyncSession and async engine. ### Install ```bash - pip install fastapi-async-sqlalchemy + pip install fastapi-async-sqlalchemy ``` diff --git a/fastapi_async_sqlalchemy/__init__.py b/fastapi_async_sqlalchemy/__init__.py index 8afd39f..963466d 100644 --- a/fastapi_async_sqlalchemy/__init__.py +++ b/fastapi_async_sqlalchemy/__init__.py @@ -1,5 +1,9 @@ -from fastapi_async_sqlalchemy.middleware import SQLAlchemyMiddleware, db +from fastapi_async_sqlalchemy.middleware import ( + SQLAlchemyMiddleware, + db, + create_middleware_and_session_proxy, +) -__all__ = ["db", "SQLAlchemyMiddleware"] +__all__ = ["db", "SQLAlchemyMiddleware", "create_middleware_and_session_proxy"] -__version__ = "0.7.0.dev3" +__version__ = "0.7.0.dev4"