Merge "Add 'FOR NO KEY UPDATE' / 'FOR KEY SHARE' support for Postgresql"

8 files changed, 141 insertions, 17 deletions

diff --git a/doc/build/changelog/changelog_11.rst b/doc/build/changelog/changelog_11.rst
index 5dee9867a..b57336c4a 100644
--- a/doc/build/changelog/changelog_11.rst
+++ b/doc/build/changelog/changelog_11.rst
@@ -110,6 +110,16 @@
         Alex Grönholm.
 
     .. change::
+        :tags: feature, postgresql
+        :pullreq: github:297
+
+        Added new parameter
+        :paramref:`.GenerativeSelect.with_for_update.key_share`, which
+        will render the ``FOR NO KEY UPDATE`` version of ``FOR UPDATE``
+        and ``FOR KEY SHARE`` instead of ``FOR SHARE``
+        on the Postgresql backend.  Pull request courtesy Sergey Skopin.
+
+    .. change::
         :tags: feature, postgresql, oracle
         :pullreq: bitbucket:86
 
diff --git a/doc/build/changelog/migration_11.rst b/doc/build/changelog/migration_11.rst
index 8c138f090..8b8075a29 100644
--- a/doc/build/changelog/migration_11.rst
+++ b/doc/build/changelog/migration_11.rst
@@ -2150,13 +2150,25 @@ should be calling upon ``sqlalchemy.dialects.postgresql``.
 Engine URLs of the form ``postgres://`` will still continue to function,
 however.
 
-Support for SKIP LOCKED
------------------------
+Support for FOR UPDATE SKIP LOCKED  / FOR NO KEY UPDATE / FOR KEY SHARE
+-----------------------------------------------------------------------
 
-The new parameter :paramref:`.GenerativeSelect.with_for_update.skip_locked`
-in both Core and ORM will generate the "SKIP LOCKED" suffix for a
-"SELECT...FOR UPDATE" or "SELECT.. FOR SHARE" query.
+The new parameters :paramref:`.GenerativeSelect.with_for_update.skip_locked`
+and :paramref:`.GenerativeSelect.with_for_update.key_share`
+in both Core and ORM apply a modification to a "SELECT...FOR UPDATE"
+or "SELECT...FOR SHARE" query on the Postgresql backend:
+
+* SELECT FOR NO KEY UPDATE::
+
+    stmt = select([table]).with_for_update(key_share=True)
+
+* SELECT FOR UPDATE SKIP LOCKED::
+
+    stmt = select([table]).with_for_update(skip_locked=True)
+
+* SELECT FOR KEY SHARE::
 
+    stmt = select([table]).with_for_update(read=True, key_share=True)
 
 Dialect Improvements and Changes - MySQL
 =============================================
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 2356458b9..d613aac92 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1170,7 +1170,12 @@ class PGCompiler(compiler.SQLCompiler):
     def for_update_clause(self, select, **kw):
 
         if select._for_update_arg.read:
-            tmp = " FOR SHARE"
+            if select._for_update_arg.key_share:
+                tmp = " FOR KEY SHARE"
+            else:
+                tmp = " FOR SHARE"
+        elif select._for_update_arg.key_share:
+            tmp = " FOR NO KEY UPDATE"
         else:
             tmp = " FOR UPDATE"
 
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 7fab33197..c1daaaf07 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -1398,7 +1398,7 @@ class Query(object):
 
     @_generative()
     def with_for_update(self, read=False, nowait=False, of=None,
-                        skip_locked=False):
+                        skip_locked=False, key_share=False):
         """return a new :class:`.Query` with the specified options for the
         ``FOR UPDATE`` clause.
 
@@ -1427,7 +1427,8 @@ class Query(object):
 
         """
         self._for_update_arg = LockmodeArg(read=read, nowait=nowait, of=of,
-                                           skip_locked=skip_locked)
+                                           skip_locked=skip_locked,
+                                           key_share=key_share)
 
     @_generative()
     def params(self, *args, **kwargs):
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index bd1d04e57..6ef327b95 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -1673,7 +1673,7 @@ class ForUpdateArg(ClauseElement):
 
     @classmethod
     def parse_legacy_select(self, arg):
-        """Parse the for_update arugment of :func:`.select`.
+        """Parse the for_update argument of :func:`.select`.
 
         :param mode: Defines the lockmode to use.
 
@@ -1723,7 +1723,9 @@ class ForUpdateArg(ClauseElement):
         if self.of is not None:
             self.of = [clone(col, **kw) for col in self.of]
 
-    def __init__(self, nowait=False, read=False, of=None, skip_locked=False):
+    def __init__(
+            self, nowait=False, read=False, of=None,
+            skip_locked=False, key_share=False):
         """Represents arguments specified to :meth:`.Select.for_update`.
 
         .. versionadded:: 0.9.0
@@ -1733,6 +1735,7 @@ class ForUpdateArg(ClauseElement):
         self.nowait = nowait
         self.read = read
         self.skip_locked = skip_locked
+        self.key_share = key_share
         if of is not None:
             self.of = [_interpret_as_column_or_from(elem)
                        for elem in util.to_list(of)]
@@ -1876,7 +1879,7 @@ class GenerativeSelect(SelectBase):
 
     @_generative
     def with_for_update(self, nowait=False, read=False, of=None,
-                        skip_locked=False):
+                        skip_locked=False, key_share=False):
         """Specify a ``FOR UPDATE`` clause for this :class:`.GenerativeSelect`.
 
         E.g.::
@@ -1917,12 +1920,16 @@ class GenerativeSelect(SelectBase):
 
          .. versionadded:: 1.1.0
 
-        .. versionadded:: 0.9.0
+        :param key_share: boolean, will render ``FOR NO KEY UPDATE``,
+         or if combined with ``read=True`` will render ``FOR KEY SHARE``,
+         on the Postgresql dialect.
 
+         .. versionadded:: 1.1.0
 
         """
         self._for_update_arg = ForUpdateArg(nowait=nowait, read=read, of=of,
-                                            skip_locked=skip_locked)
+                                            skip_locked=skip_locked,
+                                            key_share=key_share)
 
     @_generative
     def apply_labels(self):
diff --git a/test/dialect/postgresql/test_compiler.py b/test/dialect/postgresql/test_compiler.py
index c061cfaf1..c8dc9582a 100644
--- a/test/dialect/postgresql/test_compiler.py
+++ b/test/dialect/postgresql/test_compiler.py
@@ -1,7 +1,7 @@
 # coding: utf-8
 
 from sqlalchemy.testing.assertions import AssertsCompiledSQL, is_, \
-    assert_raises
+    assert_raises, assert_raises_message
 from sqlalchemy.testing import engines, fixtures
 from sqlalchemy import testing
 from sqlalchemy import Sequence, Table, Column, Integer, update, String,\
@@ -667,6 +667,58 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
             "FROM mytable WHERE mytable.myid = %(myid_1)s "
             "FOR SHARE OF mytable SKIP LOCKED")
 
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(key_share=True, nowait=True,
+                            of=[table1.c.myid, table1.c.name]),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR NO KEY UPDATE OF mytable NOWAIT")
+
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(key_share=True, skip_locked=True,
+                            of=[table1.c.myid, table1.c.name]),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR NO KEY UPDATE OF mytable SKIP LOCKED")
+
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(key_share=True,
+                            of=[table1.c.myid, table1.c.name]),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR NO KEY UPDATE OF mytable")
+
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(key_share=True),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR NO KEY UPDATE")
+
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(read=True, key_share=True),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR KEY SHARE")
+
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(read=True, key_share=True, of=table1),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR KEY SHARE OF mytable")
+
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(read=True, key_share=True, skip_locked=True),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR KEY SHARE SKIP LOCKED")
+
         ta = table1.alias()
         self.assert_compile(
             ta.select(ta.c.myid == 7).
diff --git a/test/dialect/test_oracle.py b/test/dialect/test_oracle.py
index 816741231..ed09141bb 100644
--- a/test/dialect/test_oracle.py
+++ b/test/dialect/test_oracle.py
@@ -341,6 +341,20 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
             "FROM mytable WHERE mytable.myid = :myid_1 FOR UPDATE OF "
             "mytable.myid, mytable.name SKIP LOCKED")
 
+        # key_share has no effect
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+                with_for_update(key_share=True),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = :myid_1 FOR UPDATE")
+
+        # read has no effect
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+                with_for_update(read=True, key_share=True),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = :myid_1 FOR UPDATE")
+
         ta = table1.alias()
         self.assert_compile(
             ta.select(ta.c.myid == 7).
@@ -925,7 +939,7 @@ drop synonym %(test_schema)s.local_table;
                             oracle_resolve_synonyms=True)
         self.assert_compile(parent.select(),
                 "SELECT %(test_schema)s_pt.id, "
-                "%(test_schema)s_pt.data FROM %(test_schema)s_pt" 
+                "%(test_schema)s_pt.data FROM %(test_schema)s_pt"
                  % {"test_schema": testing.config.test_schema})
         select([parent]).execute().fetchall()
 
diff --git a/test/orm/test_lockmode.py b/test/orm/test_lockmode.py
index 949fe0d81..078ffd52a 100644
--- a/test/orm/test_lockmode.py
+++ b/test/orm/test_lockmode.py
@@ -53,17 +53,18 @@ class LegacyLockModeTest(_fixtures.FixtureTest):
             sess.query(User.id).with_lockmode, 'unknown_mode'
         )
 
+
 class ForUpdateTest(_fixtures.FixtureTest):
     @classmethod
     def setup_mappers(cls):
         User, users = cls.classes.User, cls.tables.users
         mapper(User, users)
 
-    def _assert(self, read=False, nowait=False, of=None,
+    def _assert(self, read=False, nowait=False, of=None, key_share=None,
                     assert_q_of=None, assert_sel_of=None):
         User = self.classes.User
         s = Session()
-        q = s.query(User).with_for_update(read=read, nowait=nowait, of=of)
+        q = s.query(User).with_for_update(read=read, nowait=nowait, of=of, key_share=key_share)
         sel = q._compile_context().statement
 
         assert q._for_update_arg.read is read
@@ -72,9 +73,15 @@ class ForUpdateTest(_fixtures.FixtureTest):
         assert q._for_update_arg.nowait is nowait
         assert sel._for_update_arg.nowait is nowait
 
+        assert q._for_update_arg.key_share is key_share
+        assert sel._for_update_arg.key_share is key_share
+
         eq_(q._for_update_arg.of, assert_q_of)
         eq_(sel._for_update_arg.of, assert_sel_of)
 
+    def test_key_share(self):
+        self._assert(key_share=True)
+
     def test_read(self):
         self._assert(read=True)
 
@@ -172,6 +179,22 @@ class CompileTest(_fixtures.FixtureTest, AssertsCompiledSQL):
             dialect="postgresql"
         )
 
+    def test_postgres_for_no_key_update(self):
+        User = self.classes.User
+        sess = Session()
+        self.assert_compile(sess.query(User.id).with_for_update(key_share=True),
+            "SELECT users.id AS users_id FROM users FOR NO KEY UPDATE",
+            dialect="postgresql"
+        )
+
+    def test_postgres_for_no_key_nowait_update(self):
+        User = self.classes.User
+        sess = Session()
+        self.assert_compile(sess.query(User.id).with_for_update(key_share=True, nowait=True),
+            "SELECT users.id AS users_id FROM users FOR NO KEY UPDATE NOWAIT",
+            dialect="postgresql"
+        )
+
     def test_postgres_update_of_list(self):
         User = self.classes.User
         sess = Session()