allow single tables and entities for "of"

Opened up typing on :meth:`.Select.with_for_update.of` to also accept table
and mapped class arguments, as seems to be available for the MySQL dialect.

Fixes: #9174
Change-Id: I15659d7084657564bd5a2aa55ef0e4db51b91247

diff --git a/doc/build/changelog/unreleased_20/9174.rst b/doc/build/changelog/unreleased_20/9174.rst
new file mode 100644 (file)
index 0000000..3defc72
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/9174.rst
@@ -0,0 +1,6 @@
+.. change::
+    :tags: bug, typing
+    :tickets: 9174
+
+    Opened up typing on :meth:`.Select.with_for_update.of` to also accept table
+    and mapped class arguments, as seems to be available for the MySQL dialect.

diff --git a/lib/sqlalchemy/sql/_typing.py b/lib/sqlalchemy/sql/_typing.py
index da3a9ad4e7a8fb1e2d7c68278bbb439ae85f9adb..5c1a501e47e0bf3b29b030360a91393704fa29f8 100644 (file)
--- a/lib/sqlalchemy/sql/_typing.py
+++ b/lib/sqlalchemy/sql/_typing.py
@@ -186,7 +186,6 @@ overall which brings in the TextClause object also.
 
 """
 
-
 _ColumnExpressionOrLiteralArgument = Union[Any, _ColumnExpressionArgument[_T]]
 
 _ColumnExpressionOrStrLabelArgument = Union[str, _ColumnExpressionArgument[_T]]

diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index 8fb9ea0cde9a1df93aef5c92f0fd76c9152a10c1..78fc7b1b79053a87478386a34b78e48dccd1bf21 100644 (file)
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -163,6 +163,16 @@ class _JoinTargetProtocol(Protocol):
 _JoinTargetElement = Union["FromClause", _JoinTargetProtocol]
 _OnClauseElement = Union["ColumnElement[bool]", _JoinTargetProtocol]
 
+_ForUpdateOfArgument = Union[
+    # single column, Table, ORM Entity
+    Union[
+        "_ColumnExpressionArgument[Any]",
+        "_FromClauseArgument",
+    ],
+    # or sequence of single column elements
+    Sequence["_ColumnExpressionArgument[Any]"],
+]
+
 
 _SetupJoinsElement = Tuple[
     _JoinTargetElement,
@@ -3151,12 +3161,7 @@ class ForUpdateArg(ClauseElement):
         *,
         nowait: bool = False,
         read: bool = False,
-        of: Optional[
-            Union[
-                _ColumnExpressionArgument[Any],
-                Sequence[_ColumnExpressionArgument[Any]],
-            ]
-        ] = None,
+        of: Optional[_ForUpdateOfArgument] = None,
         skip_locked: bool = False,
         key_share: bool = False,
     ):
@@ -3822,12 +3827,7 @@ class GenerativeSelect(SelectBase, Generative):
         *,
         nowait: bool = False,
         read: bool = False,
-        of: Optional[
-            Union[
-                _ColumnExpressionArgument[Any],
-                Sequence[_ColumnExpressionArgument[Any]],
-            ]
-        ] = None,
+        of: Optional[_ForUpdateOfArgument] = None,
         skip_locked: bool = False,
         key_share: bool = False,
     ) -> SelfGenerativeSelect:
@@ -3860,12 +3860,12 @@ class GenerativeSelect(SelectBase, Generative):
          ``FOR SHARE`` on PostgreSQL.  On PostgreSQL, when combined with
          ``nowait``, will render ``FOR SHARE NOWAIT``.
 
-        :param of: SQL expression or list of SQL expression elements
-         (typically :class:`_schema.Column`
-         objects or a compatible expression) which
-         will render into a ``FOR UPDATE OF`` clause; supported by PostgreSQL
-         and Oracle.  May render as a table or as a column depending on
-         backend.
+        :param of: SQL expression or list of SQL expression elements,
+         (typically :class:`_schema.Column` objects or a compatible expression,
+         for some backends may also be a table expression) which will render
+         into a ``FOR UPDATE OF`` clause; supported by PostgreSQL Oracle, some
+         MySQL versions and possibly others. May render as a table or as a
+         column depending on backend.
 
         :param skip_locked: boolean, will render ``FOR UPDATE SKIP LOCKED``
          on Oracle and PostgreSQL dialects or ``FOR SHARE SKIP LOCKED`` if

diff --git a/test/ext/mypy/plain_files/common_sql_element.py b/test/ext/mypy/plain_files/common_sql_element.py
index 586a130d2596a00bc5e10f653c80d2f10ceff867..fd90e31a11c1a01da6e0529b2d332898cc342dfc 100644 (file)
--- a/test/ext/mypy/plain_files/common_sql_element.py
+++ b/test/ext/mypy/plain_files/common_sql_element.py
@@ -128,3 +128,13 @@ q1 = (
 )
 # EXPECTED_TYPE: RowReturningQuery[Tuple[int]]
 reveal_type(q1)
+
+# test 9174
+s9174_1 = select(User).with_for_update(of=User)
+s9174_2 = select(User).with_for_update(of=User.id)
+s9174_3 = select(User).with_for_update(of=[User.id, User.email])
+s9174_4 = select(user_table).with_for_update(of=user_table)
+s9174_5 = select(user_table).with_for_update(of=user_table.c.id)
+s9174_6 = select(user_table).with_for_update(
+    of=[user_table.c.id, user_table.c.email]
+)