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, 21 insertions, 8 deletions

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