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

Adds ``key_share=True`` for with_for_update().

Co-authored-by: Mike Bayer <mike_mp@zzzcomputing.com>
Change-Id: I74e0c3fcbc023e1dc98a1fa0c7db67b4c3693a31
Pull-request: https://github.com/zzzeek/sqlalchemy/pull/279

3 files changed, 93 insertions, 4 deletions

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()