Merge "Add SKIP LOCKED support for Postgresql, Oracle"

9 files changed, 94 insertions, 5 deletions

diff --git a/doc/build/changelog/changelog_11.rst b/doc/build/changelog/changelog_11.rst
index 40fa46b7c..f58b3e328 100644
--- a/doc/build/changelog/changelog_11.rst
+++ b/doc/build/changelog/changelog_11.rst
@@ -80,6 +80,16 @@
         Alex Grönholm.
 
     .. change::
+        :tags: feature, postgresql, oracle
+        :pullreq: bitbucket:86
+
+        Added new parameter
+        :paramref:`.GenerativeSelect.with_for_update.skip_locked`, which
+        will render the ``SKIP LOCKED`` phrase for a ``FOR UPDATE`` or
+        ``FOR SHARE`` lock on the Postgresql and Oracle backends.  Pull
+        request courtesy Jack Zhou.
+
+    .. change::
         :tags: change, orm
         :tickets: 3394
 
diff --git a/doc/build/changelog/migration_11.rst b/doc/build/changelog/migration_11.rst
index 7ebaf84a0..90f76fef0 100644
--- a/doc/build/changelog/migration_11.rst
+++ b/doc/build/changelog/migration_11.rst
@@ -2097,6 +2097,14 @@ should be calling upon ``sqlalchemy.dialects.postgresql``.
 Engine URLs of the form ``postgres://`` will still continue to function,
 however.
 
+Support for SKIP LOCKED
+-----------------------
+
+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.
+
+
 Dialect Improvements and Changes - MySQL
 =============================================
 
@@ -2381,3 +2389,10 @@ on this behavior, set the flag back to True.
 
 Dialect Improvements and Changes - Oracle
 =============================================
+
+Support for SKIP LOCKED
+-----------------------
+
+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.
diff --git a/lib/sqlalchemy/dialects/oracle/base.py b/lib/sqlalchemy/dialects/oracle/base.py
index 6992670c4..493bf362f 100644
--- a/lib/sqlalchemy/dialects/oracle/base.py
+++ b/lib/sqlalchemy/dialects/oracle/base.py
@@ -814,6 +814,8 @@ class OracleCompiler(compiler.SQLCompiler):
 
         if select._for_update_arg.nowait:
             tmp += " NOWAIT"
+        if select._for_update_arg.skip_locked:
+            tmp += " SKIP LOCKED"
 
         return tmp
 
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 136cb1b28..1bc4409f2 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1185,6 +1185,8 @@ class PGCompiler(compiler.SQLCompiler):
 
         if select._for_update_arg.nowait:
             tmp += " NOWAIT"
+        if select._for_update_arg.skip_locked:
+            tmp += " SKIP LOCKED"
 
         return tmp
 
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 34daa707f..7fab33197 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -1397,7 +1397,8 @@ 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,
+                        skip_locked=False):
         """return a new :class:`.Query` with the specified options for the
         ``FOR UPDATE`` clause.
 
@@ -1425,7 +1426,8 @@ 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,
+                                           skip_locked=skip_locked)
 
     @_generative()
     def params(self, *args, **kwargs):
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index e299f067e..bd1d04e57 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -1723,14 +1723,16 @@ 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, skip_locked=False):
         """Represents arguments specified to :meth:`.Select.for_update`.
 
         .. versionadded:: 0.9.0
+
         """
 
         self.nowait = nowait
         self.read = read
+        self.skip_locked = skip_locked
         if of is not None:
             self.of = [_interpret_as_column_or_from(elem)
                        for elem in util.to_list(of)]
@@ -1873,7 +1875,8 @@ 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,
+                        skip_locked=False):
         """Specify a ``FOR UPDATE`` clause for this :class:`.GenerativeSelect`.
 
         E.g.::
@@ -1908,10 +1911,18 @@ class GenerativeSelect(SelectBase):
          and Oracle.  May render as a table or as a column depending on
          backend.
 
+        :param skip_locked: boolean, will render ``FOR UPDATE SKIP LOCKED``
+         on Oracle and Postgresql dialects or ``FOR SHARE SKIP LOCKED`` if
+         ``read=True`` is also specified.
+
+         .. versionadded:: 1.1.0
+
         .. 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,
+                                            skip_locked=skip_locked)
 
     @_generative
     def apply_labels(self):
diff --git a/test/dialect/postgresql/test_compiler.py b/test/dialect/postgresql/test_compiler.py
index c20e48b01..c061cfaf1 100644
--- a/test/dialect/postgresql/test_compiler.py
+++ b/test/dialect/postgresql/test_compiler.py
@@ -606,6 +606,13 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
             "FROM mytable WHERE mytable.myid = %(myid_1)s FOR UPDATE NOWAIT")
 
         self.assert_compile(
+            table1.select(table1.c.myid == 7).
+            with_for_update(skip_locked=True),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR UPDATE SKIP LOCKED")
+
+        self.assert_compile(
             table1.select(table1.c.myid == 7).with_for_update(read=True),
             "SELECT mytable.myid, mytable.name, mytable.description "
             "FROM mytable WHERE mytable.myid = %(myid_1)s FOR SHARE")
@@ -618,6 +625,13 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
 
         self.assert_compile(
             table1.select(table1.c.myid == 7).
+            with_for_update(read=True, skip_locked=True),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = %(myid_1)s "
+            "FOR SHARE SKIP LOCKED")
+
+        self.assert_compile(
+            table1.select(table1.c.myid == 7).
             with_for_update(of=table1.c.myid),
             "SELECT mytable.myid, mytable.name, mytable.description "
             "FROM mytable WHERE mytable.myid = %(myid_1)s "
@@ -645,6 +659,14 @@ 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(read=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 SHARE OF mytable 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 1f3e63040..816741231 100644
--- a/test/dialect/test_oracle.py
+++ b/test/dialect/test_oracle.py
@@ -334,6 +334,13 @@ 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(skip_locked=True, of=[table1.c.myid, table1.c.name]),
+            "SELECT mytable.myid, mytable.name, mytable.description "
+            "FROM mytable WHERE mytable.myid = :myid_1 FOR UPDATE OF "
+            "mytable.myid, mytable.name SKIP LOCKED")
+
         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..949fe0d81 100644
--- a/test/orm/test_lockmode.py
+++ b/test/orm/test_lockmode.py
@@ -181,6 +181,15 @@ class CompileTest(_fixtures.FixtureTest, AssertsCompiledSQL):
             dialect="postgresql"
         )
 
+    def test_postgres_update_skip_locked(self):
+        User = self.classes.User
+        sess = Session()
+        self.assert_compile(sess.query(User.id).
+                with_for_update(skip_locked=True),
+            "SELECT users.id AS users_id FROM users FOR UPDATE SKIP LOCKED",
+            dialect="postgresql"
+        )
+
 
     def test_oracle_update(self):
         User = self.classes.User
@@ -190,6 +199,15 @@ class CompileTest(_fixtures.FixtureTest, AssertsCompiledSQL):
             dialect="oracle"
         )
 
+    def test_oracle_update_skip_locked(self):
+        User = self.classes.User
+        sess = Session()
+        self.assert_compile(sess.query(User.id)
+                .with_for_update(skip_locked=True),
+            "SELECT users.id AS users_id FROM users FOR UPDATE SKIP LOCKED",
+            dialect="oracle"
+        )
+
     def test_mysql_read(self):
         User = self.classes.User
         sess = Session()