Add 'FOR NO KEY UPDATE' supportpr/279

6 files changed, 123 insertions, 16 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 9d019b56e..5407aa4be 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1159,6 +1159,8 @@ class PGCompiler(compiler.SQLCompiler):
 
         if select._for_update_arg.read:
             tmp = " FOR SHARE"
+        elif select._for_update_arg.no_key:
+            tmp = " FOR NO KEY UPDATE"
         else:
             tmp = " FOR UPDATE"
 
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 34daa707f..2b1a88582 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -1388,6 +1388,12 @@ class Query(object):
          * ``'read'`` - translates to ``LOCK IN SHARE MODE`` (for MySQL),
            and ``FOR SHARE`` (for PostgreSQL)
 
+         * ``'no_key'`` - translates to ``FOR NO KEY UPDATE`` (for PostgreSQL),
+           and ``FOR UPDATE`` on other backends
+
+         * ``'no_key_nowait'`` - translates to ``FOR NO KEY UPDATE NOWAIT`` (for PostgreSQL),
+           and ``FOR UPDATE`` on other backends
+
         .. seealso::
 
             :meth:`.Query.with_for_update` - improved API for
@@ -1397,7 +1403,7 @@ class Query(object):
         self._for_update_arg = LockmodeArg.parse_legacy_query(mode)
 
     @_generative()
-    def with_for_update(self, read=False, nowait=False, of=None):
+    def with_for_update(self, read=False, nowait=False, of=None, no_key=None):
         """return a new :class:`.Query` with the specified options for the
         ``FOR UPDATE`` clause.
 
@@ -1411,10 +1417,12 @@ class Query(object):
         E.g.::
 
             q = sess.query(User).with_for_update(nowait=True, of=User)
+            q = sess.query(User).with_for_update(nowait=True, no_key=True)
 
-        The above query on a Postgresql backend will render like::
+        The above queries on a Postgresql backend will render like, accordingly::
 
             SELECT users.id AS users_id FROM users FOR UPDATE OF users NOWAIT
+            SELECT users.id AS users_id FROM users FOR NO KEY UPDATE NOWAIT
 
         .. versionadded:: 0.9.0 :meth:`.Query.with_for_update` supersedes
            the :meth:`.Query.with_lockmode` method.
@@ -1425,7 +1433,7 @@ class Query(object):
             full argument and behavioral description.
 
         """
-        self._for_update_arg = LockmodeArg(read=read, nowait=nowait, of=of)
+        self._for_update_arg = LockmodeArg(read=read, nowait=nowait, of=of, no_key=no_key)
 
     @_generative()
     def params(self, *args, **kwargs):
@@ -3411,6 +3419,7 @@ class LockmodeArg(ForUpdateArg):
         if mode in (None, False):
             return None
 
+        no_key = nowait = read = False
         if mode == "read":
             read = True
             nowait = False
@@ -3419,11 +3428,15 @@ class LockmodeArg(ForUpdateArg):
         elif mode == "update_nowait":
             nowait = True
             read = False
+        elif mode == "no_key":
+            no_key = True
+        elif mode == "no_key_nowait":
+            no_key = nowait = True
         else:
             raise sa_exc.ArgumentError(
                 "Unknown with_lockmode argument: %r" % mode)
 
-        return LockmodeArg(read=read, nowait=nowait)
+        return LockmodeArg(read=read, nowait=nowait, no_key=no_key)
 
 
 class _QueryEntity(object):
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index e299f067e..38c8a1453 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.
 
@@ -1688,6 +1688,12 @@ class ForUpdateArg(ClauseElement):
             ``'read'`` - translates to ``LOCK IN SHARE MODE`` (for MySQL),
             and ``FOR SHARE`` (for PostgreSQL)
 
+            ``'no_key'`` - translates to ``FOR NO KEY UPDATE`` (for PostgreSQL),
+            and ``FOR UPDATE`` on other backends
+
+            ``'no_key_nowait'`` - translates to ``FOR NO KEY UPDATE NOWAIT`` (for PostgreSQL),
+            and ``FOR UPDATE`` on other backends
+
             ``'read_nowait'`` - translates to ``FOR SHARE NOWAIT``
             (supported by PostgreSQL). ``FOR SHARE`` and
             ``FOR SHARE NOWAIT`` (PostgreSQL).
@@ -1696,17 +1702,21 @@ class ForUpdateArg(ClauseElement):
         if arg in (None, False):
             return None
 
-        nowait = read = False
+        nowait = read = no_key = False
         if arg == 'nowait':
             nowait = True
         elif arg == 'read':
             read = True
         elif arg == 'read_nowait':
             read = nowait = True
+        elif arg == 'no_key':
+            no_key = True
+        elif arg == 'no_key_nowait':
+            no_key = nowait = True
         elif arg is not True:
             raise exc.ArgumentError("Unknown for_update argument: %r" % arg)
 
-        return ForUpdateArg(read=read, nowait=nowait)
+        return ForUpdateArg(read=read, nowait=nowait, no_key=no_key)
 
     @property
     def legacy_for_update_value(self):
@@ -1723,7 +1733,7 @@ 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):
+    def __init__(self, nowait=False, read=False, of=None, no_key=None):
         """Represents arguments specified to :meth:`.Select.for_update`.
 
         .. versionadded:: 0.9.0
@@ -1731,12 +1741,16 @@ class ForUpdateArg(ClauseElement):
 
         self.nowait = nowait
         self.read = read
+        self.no_key = no_key
         if of is not None:
             self.of = [_interpret_as_column_or_from(elem)
                        for elem in util.to_list(of)]
         else:
             self.of = None
 
+        if self.read and self.no_key:
+            raise exc.ArgumentError("Bad argument values: 'no_key' and 'read' cannot be True together.")
+
 
 class SelectBase(HasCTE, Executable, FromClause):
     """Base class for SELECT statements.
@@ -1873,19 +1887,23 @@ class GenerativeSelect(SelectBase):
         self._for_update_arg = ForUpdateArg.parse_legacy_select(value)
 
     @_generative
-    def with_for_update(self, nowait=False, read=False, of=None):
+    def with_for_update(self, nowait=False, read=False, of=None, no_key=None):
         """Specify a ``FOR UPDATE`` clause for this :class:`.GenerativeSelect`.
 
         E.g.::
 
-            stmt = select([table]).with_for_update(nowait=True)
+            stmt = select([table]).with_for_update(nowait=True, no_key=True)
 
-        On a database like Postgresql or Oracle, the above would render a
+        On a database like Postgresql, the above would render a
         statement like::
 
+            SELECT table.a, table.b FROM table FOR NO KEY UPDATE NOWAIT
+
+        on Oracle::
+
             SELECT table.a, table.b FROM table FOR UPDATE NOWAIT
 
-        on other backends, the ``nowait`` option is ignored and instead
+        on other backends, the ``nowait`` and ``no_key`` options is ignored and instead
         would produce::
 
             SELECT table.a, table.b FROM table FOR UPDATE
@@ -1911,7 +1929,7 @@ class GenerativeSelect(SelectBase):
         .. versionadded:: 0.9.0
 
         """
-        self._for_update_arg = ForUpdateArg(nowait=nowait, read=read, of=of)
+        self._for_update_arg = ForUpdateArg(nowait=nowait, read=read, of=of, no_key=no_key)
 
     @_generative
     def apply_labels(self):
diff --git a/test/dialect/postgresql/test_compiler.py b/test/dialect/postgresql/test_compiler.py
index 87e48d3f2..6ff1102e6 100644
--- a/test/dialect/postgresql/test_compiler.py
+++ b/test/dialect/postgresql/test_compiler.py
@@ -602,6 +602,43 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
             "FROM mytable WHERE mytable.myid = %(myid_1)s "
             "FOR SHARE OF mytable NOWAIT")
 
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(no_key=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(no_key=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(no_key=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(no_key=True),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR NO KEY UPDATE")
+
+        assert_raises(
+            exc.ArgumentError,
+            table1.select(table1.c.myid == 7).with_for_update,
+            no_key=True,
+            read=True
+        )
+
         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 1f3e63040..b715667f9 100644
--- a/test/dialect/test_oracle.py
+++ b/test/dialect/test_oracle.py
@@ -334,6 +334,12 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
             "FROM mytable WHERE mytable.myid = :myid_1 FOR UPDATE OF "
             "mytable.myid, mytable.name NOWAIT")
 
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
+                with_for_update(no_key=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).
diff --git a/test/orm/test_lockmode.py b/test/orm/test_lockmode.py
index fc473a329..fa9588ac7 100644
--- a/test/orm/test_lockmode.py
+++ b/test/orm/test_lockmode.py
@@ -16,7 +16,7 @@ class LegacyLockModeTest(_fixtures.FixtureTest):
         User, users = cls.classes.User, cls.tables.users
         mapper(User, users)
 
-    def _assert_legacy(self, arg, read=False, nowait=False):
+    def _assert_legacy(self, arg, read=False, nowait=False, no_key=False):
         User = self.classes.User
         s = Session()
         q = s.query(User).with_lockmode(arg)
@@ -29,9 +29,11 @@ class LegacyLockModeTest(_fixtures.FixtureTest):
 
         assert q._for_update_arg.read is read
         assert q._for_update_arg.nowait is nowait
+        assert q._for_update_arg.no_key is no_key
 
         assert sel._for_update_arg.read is read
         assert sel._for_update_arg.nowait is nowait
+        assert sel._for_update_arg.no_key is no_key
 
     def test_false_legacy(self):
         self._assert_legacy(None)
@@ -42,6 +44,12 @@ class LegacyLockModeTest(_fixtures.FixtureTest):
     def test_nowait_legacy(self):
         self._assert_legacy("update_nowait", nowait=True)
 
+    def test_no_key_nowait_legacy(self):
+        self._assert_legacy("no_key_nowait", nowait=True, no_key=True)
+
+    def test_no_key_legacy(self):
+        self._assert_legacy("no_key", no_key=True)
+
     def test_read_legacy(self):
         self._assert_legacy("read", read=True)
 
@@ -53,17 +61,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, no_key=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, no_key=no_key)
         sel = q._compile_context().statement
 
         assert q._for_update_arg.read is read
@@ -72,9 +81,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.no_key is no_key
+        assert sel._for_update_arg.no_key is no_key
+
         eq_(q._for_update_arg.of, assert_q_of)
         eq_(sel._for_update_arg.of, assert_sel_of)
 
+    def test_no_key(self):
+        self._assert(no_key=True)
+
     def test_read(self):
         self._assert(read=True)
 
@@ -172,6 +187,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(no_key=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(no_key=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()