From cae662a6383d3ae8f3673c70c3118ea3a1a1606e Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Sat, 28 Jan 2023 23:41:42 -0500 Subject: [PATCH] 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 --- doc/build/changelog/unreleased_20/9174.rst | 6 ++++ lib/sqlalchemy/sql/_typing.py | 1 - lib/sqlalchemy/sql/selectable.py | 36 +++++++++---------- .../mypy/plain_files/common_sql_element.py | 10 ++++++ 4 files changed, 34 insertions(+), 19 deletions(-) create mode 100644 doc/build/changelog/unreleased_20/9174.rst diff --git a/doc/build/changelog/unreleased_20/9174.rst b/doc/build/changelog/unreleased_20/9174.rst new file mode 100644 index 0000000000..3defc72c46 --- /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 da3a9ad4e7..5c1a501e47 100644 --- 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 8fb9ea0cde..78fc7b1b79 100644 --- 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 586a130d25..fd90e31a11 100644 --- 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] +) -- 2.47.2