fix: correct race condition in with_for_update

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
       - id: unasyncd
         additional_dependencies: ["ruff"]
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.14.2"
+    rev: "v0.14.5"
     hooks:
       # Run the linter.
       - id: ruff

advanced_alchemy/repository/_async.py

@@ -39,7 +39,7 @@
 from sqlalchemy.orm.strategy_options import _AbstractLoad  # pyright: ignore[reportPrivateUsage]
 from sqlalchemy.sql import ColumnElement
 from sqlalchemy.sql.dml import ReturningDelete, ReturningUpdate
-from sqlalchemy.sql.selectable import ForUpdateParameter
+from sqlalchemy.sql.selectable import ForUpdateArg, ForUpdateParameter
 
 from advanced_alchemy.exceptions import ErrorMessages, NotFoundError, RepositoryError, wrap_sqlalchemy_exception
 from advanced_alchemy.filters import StatementFilter, StatementTypeT
@@ -202,6 +202,7 @@ async def get(
         error_messages: Optional[Union[ErrorMessages, EmptyType]] = Empty,
         load: Optional[LoadSpec] = None,
         execution_options: Optional[dict[str, Any]] = None,
+        with_for_update: ForUpdateParameter = None,
     ) -> ModelT: ...
 
     async def get_one(
@@ -1033,6 +1034,31 @@ def _get_base_stmt(
             statement = cast("StatementTypeT", statement.execution_options(**execution_options))
         return statement
 
+    def _apply_for_update_options(
+        self,
+        statement: Select[tuple[ModelT]],
+        with_for_update: ForUpdateParameter,
+    ) -> Select[tuple[ModelT]]:
+        """Apply FOR UPDATE options to a SELECT statement when requested."""
+
+        if with_for_update in (None, False):
+            return statement
+        if with_for_update is True:
+            return statement.with_for_update()
+        if isinstance(with_for_update, ForUpdateArg):
+            with_for_update_kwargs: dict[str, Any] = {
+                "nowait": with_for_update.nowait,
+                "read": with_for_update.read,
+                "skip_locked": with_for_update.skip_locked,
+                "key_share": with_for_update.key_share,
+            }
+            if getattr(with_for_update, "of", None):
+                with_for_update_kwargs["of"] = with_for_update.of
+            return statement.with_for_update(**with_for_update_kwargs)
+        if isinstance(with_for_update, dict):  # pyright: ignore
+            return statement.with_for_update(**with_for_update)
+        return statement
+
     def _get_delete_many_statement(
         self,
         *,
@@ -1071,6 +1097,7 @@ async def get(
         load: Optional[LoadSpec] = None,
         execution_options: Optional[dict[str, Any]] = None,
         uniquify: Optional[bool] = None,
+        with_for_update: ForUpdateParameter = None,
     ) -> ModelT:
         """Get instance identified by `item_id`.
 
@@ -1085,6 +1112,7 @@ async def get(
             load: Set relationships to be loaded
             execution_options: Set default execution options
             uniquify: Optionally apply the ``unique()`` method to results before returning.
+            with_for_update: Optional FOR UPDATE clause / parameters to apply to the SELECT statement.
 
         Returns:
             The retrieved instance.
@@ -1107,6 +1135,7 @@ async def get(
                 execution_options=execution_options,
             )
             statement = self._filter_select_by_kwargs(statement, [(id_attribute, item_id)])
+            statement = self._apply_for_update_options(statement, with_for_update)
             instance = (await self._execute(statement, uniquify=loader_options_have_wildcard)).scalar_one_or_none()
             instance = self.check_not_found(instance)
             self._expunge(instance, auto_expunge=auto_expunge)
@@ -1486,7 +1515,11 @@ async def update(
                 id_attribute=id_attribute,
             )
             existing_instance = await self.get(
-                item_id, id_attribute=id_attribute, load=load, execution_options=execution_options
+                item_id,
+                id_attribute=id_attribute,
+                load=load,
+                execution_options=execution_options,
+                with_for_update=with_for_update,
             )
             mapper = None
             with (

advanced_alchemy/repository/_sync.py

@@ -40,7 +40,7 @@
 from sqlalchemy.orm.strategy_options import _AbstractLoad  # pyright: ignore[reportPrivateUsage]
 from sqlalchemy.sql import ColumnElement
 from sqlalchemy.sql.dml import ReturningDelete, ReturningUpdate
-from sqlalchemy.sql.selectable import ForUpdateParameter
+from sqlalchemy.sql.selectable import ForUpdateArg, ForUpdateParameter
 
 from advanced_alchemy.exceptions import ErrorMessages, NotFoundError, RepositoryError, wrap_sqlalchemy_exception
 from advanced_alchemy.filters import StatementFilter, StatementTypeT
@@ -203,6 +203,7 @@ def get(
         error_messages: Optional[Union[ErrorMessages, EmptyType]] = Empty,
         load: Optional[LoadSpec] = None,
         execution_options: Optional[dict[str, Any]] = None,
+        with_for_update: ForUpdateParameter = None,
     ) -> ModelT: ...
 
     def get_one(
@@ -1034,6 +1035,31 @@ def _get_base_stmt(
             statement = cast("StatementTypeT", statement.execution_options(**execution_options))
         return statement
 
+    def _apply_for_update_options(
+        self,
+        statement: Select[tuple[ModelT]],
+        with_for_update: ForUpdateParameter,
+    ) -> Select[tuple[ModelT]]:
+        """Apply FOR UPDATE options to a SELECT statement when requested."""
+
+        if with_for_update in (None, False):
+            return statement
+        if with_for_update is True:
+            return statement.with_for_update()
+        if isinstance(with_for_update, ForUpdateArg):
+            with_for_update_kwargs: dict[str, Any] = {
+                "nowait": with_for_update.nowait,
+                "read": with_for_update.read,
+                "skip_locked": with_for_update.skip_locked,
+                "key_share": with_for_update.key_share,
+            }
+            if getattr(with_for_update, "of", None):
+                with_for_update_kwargs["of"] = with_for_update.of
+            return statement.with_for_update(**with_for_update_kwargs)
+        if isinstance(with_for_update, dict):  # pyright: ignore
+            return statement.with_for_update(**with_for_update)
+        return statement
+
     def _get_delete_many_statement(
         self,
         *,
@@ -1072,6 +1098,7 @@ def get(
         load: Optional[LoadSpec] = None,
         execution_options: Optional[dict[str, Any]] = None,
         uniquify: Optional[bool] = None,
+        with_for_update: ForUpdateParameter = None,
     ) -> ModelT:
         """Get instance identified by `item_id`.
 
@@ -1086,6 +1113,7 @@ def get(
             load: Set relationships to be loaded
             execution_options: Set default execution options
             uniquify: Optionally apply the ``unique()`` method to results before returning.
+            with_for_update: Optional FOR UPDATE clause / parameters to apply to the SELECT statement.
 
         Returns:
             The retrieved instance.
@@ -1108,6 +1136,7 @@ def get(
                 execution_options=execution_options,
             )
             statement = self._filter_select_by_kwargs(statement, [(id_attribute, item_id)])
+            statement = self._apply_for_update_options(statement, with_for_update)
             instance = (self._execute(statement, uniquify=loader_options_have_wildcard)).scalar_one_or_none()
             instance = self.check_not_found(instance)
             self._expunge(instance, auto_expunge=auto_expunge)
@@ -1487,7 +1516,11 @@ def update(
                 id_attribute=id_attribute,
             )
             existing_instance = self.get(
-                item_id, id_attribute=id_attribute, load=load, execution_options=execution_options
+                item_id,
+                id_attribute=id_attribute,
+                load=load,
+                execution_options=execution_options,
+                with_for_update=with_for_update,
             )
             mapper = None
             with (

advanced_alchemy/repository/memory/_async.py

@@ -454,6 +454,7 @@ async def get(
         load: Optional[LoadSpec] = None,
         execution_options: Optional[dict[str, Any]] = None,
         uniquify: Optional[bool] = None,
+        with_for_update: ForUpdateParameter = None,
     ) -> ModelT:
         return self._find_or_raise_not_found(item_id)
 

advanced_alchemy/repository/memory/_sync.py

@@ -455,6 +455,7 @@ def get(
         load: Optional[LoadSpec] = None,
         execution_options: Optional[dict[str, Any]] = None,
         uniquify: Optional[bool] = None,
+        with_for_update: ForUpdateParameter = None,
     ) -> ModelT:
         return self._find_or_raise_not_found(item_id)
 

advanced_alchemy/service/_async.py

@@ -746,8 +746,12 @@ async def update(
         if item_id is not None:
             # When item_id is provided, update existing instance rather than replacing it
             # This preserves relationships and database-managed fields
-            existing_instance = await self.repository.get(
-                item_id, id_attribute=id_attribute, load=load, execution_options=execution_options
+            existing_instance: ModelT = await self.repository.get(
+                item_id,
+                id_attribute=id_attribute,
+                load=load,
+                execution_options=execution_options,
+                with_for_update=with_for_update,
             )
 
             # Extract attributes from converted model to update existing instance

advanced_alchemy/service/_sync.py

@@ -745,8 +745,12 @@ def update(
         if item_id is not None:
             # When item_id is provided, update existing instance rather than replacing it
             # This preserves relationships and database-managed fields
-            existing_instance = self.repository.get(
-                item_id, id_attribute=id_attribute, load=load, execution_options=execution_options
+            existing_instance: ModelT = self.repository.get(
+                item_id,
+                id_attribute=id_attribute,
+                load=load,
+                execution_options=execution_options,
+                with_for_update=with_for_update,
             )
 
             # Extract attributes from converted model to update existing instance

pyproject.toml

@@ -462,6 +462,16 @@ module = [
   "dishka.*",
 ]
 
+[[tool.mypy.overrides]]
+follow_imports = "skip"
+ignore_missing_imports = true
+module = [
+  "pytest",
+  "pytest.*",
+  "_pytest",
+  "_pytest.*",
+]
+
 [[tool.mypy.overrides]]
 module = "advanced_alchemy._serialization"
 warn_unused_ignores = false
@@ -506,6 +516,10 @@ module = "examples.flask.*"
 disable_error_code = "unreachable"
 module = "tests.integration.test_repository"
 
+[[tool.mypy.overrides]]
+module = "tests.unit.test_exceptions"
+warn_unreachable = false
+
 
 [tool.pyright]
 disableBytesTypePromotions = true

tests/unit/test_extensions/test_flask.py

@@ -4,6 +4,7 @@
 
 from collections.abc import Generator, Sequence
 from pathlib import Path
+from typing import cast
 
 import pytest
 from flask import Flask, Response
@@ -74,7 +75,7 @@ class Repo(SQLAlchemyAsyncRepository[User]):
 
 @pytest.fixture(scope="session")
 def tmp_path_session(tmp_path_factory: pytest.TempPathFactory) -> Path:
-    return tmp_path_factory.mktemp("test_extensions_flask")
+    return cast("Path", tmp_path_factory.mktemp("test_extensions_flask"))
 
 
 @pytest.fixture(scope="session")

tests/unit/test_repository.py

@@ -17,6 +17,7 @@
 from sqlalchemy.exc import InvalidRequestError, SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import InstrumentedAttribute, Mapped, Session, mapped_column
+from sqlalchemy.sql.selectable import ForUpdateArg
 from sqlalchemy.types import TypeEngine
 
 from advanced_alchemy import base
@@ -378,6 +379,89 @@ async def test_sqlalchemy_repo_get_member(
     mock_repo.session.commit.assert_not_called()  # pyright: ignore[reportFunctionMemberAccess]
 
 
+async def test_sqlalchemy_repo_get_with_for_update(
+    mock_repo: SQLAlchemyAsyncRepository[Any],
+    mocker: MockerFixture,
+) -> None:
+    """Ensure FOR UPDATE options are applied when requested."""
+
+    statement = MagicMock()
+    statement.options.return_value = statement
+    statement.execution_options.return_value = statement
+    statement.with_for_update.return_value = statement
+    mock_repo.statement = statement
+
+    mocker.patch.object(mock_repo, "_get_loader_options", return_value=([], False))
+    mocker.patch.object(mock_repo, "_get_base_stmt", return_value=statement)
+    mocker.patch.object(mock_repo, "_apply_filters", return_value=statement)
+    mocker.patch.object(mock_repo, "_filter_select_by_kwargs", return_value=statement)
+    execute_result = MagicMock()
+    execute_result.scalar_one_or_none.return_value = MagicMock()
+    execute = mocker.patch.object(mock_repo, "_execute", return_value=execute_result)
+
+    instance = await maybe_async(mock_repo.get("instance-id", with_for_update=True))
+
+    assert instance is execute_result.scalar_one_or_none.return_value
+    statement.with_for_update.assert_called_once_with()
+    execute.assert_called_once_with(statement, uniquify=False)
+
+
+async def test_sqlalchemy_repo_get_with_for_update_dict(
+    mock_repo: SQLAlchemyAsyncRepository[Any],
+    mocker: MockerFixture,
+) -> None:
+    statement = MagicMock()
+    statement.options.return_value = statement
+    statement.execution_options.return_value = statement
+    statement.with_for_update.return_value = statement
+    mock_repo.statement = statement
+
+    mocker.patch.object(mock_repo, "_get_loader_options", return_value=([], False))
+    mocker.patch.object(mock_repo, "_get_base_stmt", return_value=statement)
+    mocker.patch.object(mock_repo, "_apply_filters", return_value=statement)
+    mocker.patch.object(mock_repo, "_filter_select_by_kwargs", return_value=statement)
+    execute_result = MagicMock()
+    execute_result.scalar_one_or_none.return_value = MagicMock()
+    mocker.patch.object(mock_repo, "_execute", return_value=execute_result)
+
+    await maybe_async(
+        mock_repo.get(
+            "instance-id",
+            with_for_update={"nowait": True, "read": False},
+        )
+    )
+
+    statement.with_for_update.assert_called_once_with(nowait=True, read=False)
+
+
+async def test_sqlalchemy_repo_get_with_for_update_arg(
+    mock_repo: SQLAlchemyAsyncRepository[Any],
+    mocker: MockerFixture,
+) -> None:
+    statement = MagicMock()
+    statement.options.return_value = statement
+    statement.execution_options.return_value = statement
+    statement.with_for_update.return_value = statement
+    mock_repo.statement = statement
+
+    mocker.patch.object(mock_repo, "_get_loader_options", return_value=([], False))
+    mocker.patch.object(mock_repo, "_get_base_stmt", return_value=statement)
+    mocker.patch.object(mock_repo, "_apply_filters", return_value=statement)
+    mocker.patch.object(mock_repo, "_filter_select_by_kwargs", return_value=statement)
+    execute_result = MagicMock()
+    execute_result.scalar_one_or_none.return_value = MagicMock()
+    mocker.patch.object(mock_repo, "_execute", return_value=execute_result)
+
+    await maybe_async(
+        mock_repo.get(
+            "instance-id",
+            with_for_update=ForUpdateArg(nowait=True, key_share=True),
+        )
+    )
+
+    statement.with_for_update.assert_called_once_with(nowait=True, read=False, skip_locked=False, key_share=True)
+
+
 async def test_sqlalchemy_repo_get_one_member(
     mock_repo: SQLAlchemyAsyncRepository[Any],
     monkeypatch: MonkeyPatch,