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

1 files changed, 53 insertions, 1 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).