add psql FOR UPDATE OF functionality

5 files changed, 34 insertions, 1 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index e1dc4af71..19d7c81fa 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1015,6 +1015,8 @@ class PGCompiler(compiler.SQLCompiler):
 
     def for_update_clause(self, select):
         if select.for_update == 'nowait':
+            if select.for_update_of is not None:
+                return " FOR UPDATE OF " + select.for_update_of + " NOWAIT"
             return " FOR UPDATE NOWAIT"
         elif select.for_update == 'read':
             return " FOR SHARE"
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index c9e7d444b..db688955d 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -70,6 +70,7 @@ class Query(object):
     _criterion = None
     _yield_per = None
     _lockmode = None
+    _lockmode_of = None
     _order_by = False
     _group_by = False
     _having = None
@@ -1124,7 +1125,7 @@ class Query(object):
         self._execution_options = self._execution_options.union(kwargs)
 
     @_generative()
-    def with_lockmode(self, mode):
+    def with_lockmode(self, mode, of=None):
         """Return a new Query object with the specified locking mode.
 
         :param mode: a string representing the desired locking mode. A
@@ -1148,9 +1149,13 @@ class Query(object):
 
             .. versionadded:: 0.7.7
                 ``FOR SHARE`` and ``FOR SHARE NOWAIT`` (PostgreSQL).
+        :param of: a table descriptor representing the optional OF
+            part of the clause. This passes ``for_update_of=table'``
+            which translates to ``FOR UPDATE OF table [NOWAIT]``.
         """
 
         self._lockmode = mode
+        self._lockmode_of = of
 
     @_generative()
     def params(self, *args, **kwargs):
@@ -2705,6 +2710,9 @@ class Query(object):
             except KeyError:
                 raise sa_exc.ArgumentError(
                                 "Unknown lockmode %r" % self._lockmode)
+            if self._lockmode_of is not None:
+                context.for_update_of = self._lockmode_of
+
         for entity in self._entities:
             entity.setup_context(self, context)
 
@@ -2789,6 +2797,7 @@ class Query(object):
         statement = sql.select(
                             [inner] + context.secondary_columns,
                             for_update=context.for_update,
+                            for_update_of=context.for_update_of,
                             use_labels=context.labels)
 
         from_clause = inner
@@ -2834,6 +2843,7 @@ class Query(object):
                         from_obj=context.froms,
                         use_labels=context.labels,
                         for_update=context.for_update,
+                        for_update_of=context.for_update_of,
                         order_by=context.order_by,
                         **self._select_args
                     )
@@ -3415,6 +3425,7 @@ class QueryContext(object):
     adapter = None
     froms = ()
     for_update = False
+    for_update_of = None
 
     def __init__(self, query):
 
diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index 4f3dbba36..51ec0d9eb 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -1571,6 +1571,8 @@ class SQLCompiler(Compiled):
 
     def for_update_clause(self, select):
         if select.for_update:
+            if select.for_update_of is not None:
+                return " FOR UPDATE OF " + select.for_update_of
             return " FOR UPDATE"
         else:
             return ""
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index 550e250f1..8ad238ca3 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -1162,6 +1162,7 @@ class SelectBase(Executable, FromClause):
     def __init__(self,
             use_labels=False,
             for_update=False,
+            for_update_of=None,
             limit=None,
             offset=None,
             order_by=None,
@@ -1170,6 +1171,7 @@ class SelectBase(Executable, FromClause):
             autocommit=None):
         self.use_labels = use_labels
         self.for_update = for_update
+        self.for_update_of = for_update_of
         if autocommit is not None:
             util.warn_deprecated('autocommit on select() is '
                                  'deprecated.  Use .execution_options(a'
diff --git a/test/orm/test_lockmode.py b/test/orm/test_lockmode.py
index 0fe82f394..a16a545ba 100644
--- a/test/orm/test_lockmode.py
+++ b/test/orm/test_lockmode.py
@@ -73,6 +73,14 @@ class LockModeTest(_fixtures.FixtureTest, AssertsCompiledSQL):
             dialect=postgresql.dialect()
         )
 
+    def test_postgres_update_of(self):
+        User = self.classes.User
+        sess = Session()
+        self.assert_compile(sess.query(User.id).with_lockmode('update', of='users'),
+            "SELECT users.id AS users_id FROM users FOR UPDATE OF users",
+            dialect=postgresql.dialect()
+        )
+
     def test_postgres_update_nowait(self):
         User = self.classes.User
         sess = Session()
@@ -81,6 +89,14 @@ class LockModeTest(_fixtures.FixtureTest, AssertsCompiledSQL):
             dialect=postgresql.dialect()
         )
 
+    def test_postgres_update_nowait_of(self):
+        User = self.classes.User
+        sess = Session()
+        self.assert_compile(sess.query(User.id).with_lockmode('update_nowait', of='users'),
+            "SELECT users.id AS users_id FROM users FOR UPDATE OF users NOWAIT",
+            dialect=postgresql.dialect()
+        )
+
     def test_oracle_update(self):
         User = self.classes.User
         sess = Session()