Adapt "FOR UPDATE OF" with Oracle limit/offset

This modifies the Oracle ROWNUM limit/offset approach
to accommodate for the "OF" clause in a "FOR UPDATE"
phrase.  The column expressions must be added to the selected
subquery if necessary and adapted on the outside.

Change-Id: Ia71b5fc4df6d326e73863f8ae9f96e8f1a5acfc1
Fixes: #3741

3 files changed, 106 insertions, 3 deletions

diff --git a/doc/build/changelog/changelog_10.rst b/doc/build/changelog/changelog_10.rst
index 7240e8e68..724792174 100644
--- a/doc/build/changelog/changelog_10.rst
+++ b/doc/build/changelog/changelog_10.rst
@@ -29,6 +29,17 @@
         rather than SERIAL.
 
     .. change::
+        :tags: bug, oracle
+        :tickets: 3741
+        :versions: 1.1.0b3
+
+        Fixed bug in :paramref:`.Select.with_for_update.of`, where the Oracle
+        "rownum" approach to LIMIT/OFFSET would fail to accomodate for the
+        expressions inside the "OF" clause, which must be stated at the topmost
+        level referring to expression within the subquery.  The expressions are
+        now added to the subquery if needed.
+
+    .. change::
         :tags: bug, sql
         :tickets: 3735
         :versions: 1.1.0b2
diff --git a/lib/sqlalchemy/dialects/oracle/base.py b/lib/sqlalchemy/dialects/oracle/base.py
index 3af308cbb..a68e2d7ca 100644
--- a/lib/sqlalchemy/dialects/oracle/base.py
+++ b/lib/sqlalchemy/dialects/oracle/base.py
@@ -285,7 +285,7 @@ import re
 
 from sqlalchemy import util, sql
 from sqlalchemy.engine import default, reflection
-from sqlalchemy.sql import compiler, visitors, expression
+from sqlalchemy.sql import compiler, visitors, expression, util as sql_util
 from sqlalchemy.sql import operators as sql_operators
 from sqlalchemy.sql.elements import quoted_name
 from sqlalchemy import types as sqltypes, schema as sa_schema
@@ -754,6 +754,20 @@ class OracleCompiler(compiler.SQLCompiler):
                 limitselect._oracle_visit = True
                 limitselect._is_wrapper = True
 
+                # add expressions to accomodate FOR UPDATE OF
+                for_update = select._for_update_arg
+                if for_update is not None and for_update.of:
+                    for_update = for_update._clone()
+                    for_update._copy_internals()
+
+                    for elem in for_update.of:
+                        select.append_column(elem)
+
+                    adapter = sql_util.ClauseAdapter(select)
+                    for_update.of = [
+                        adapter.traverse(elem)
+                        for elem in for_update.of]
+
                 # If needed, add the limiting clause
                 if limit_clause is not None:
                     if not self.dialect.use_binds_for_limits:
@@ -773,7 +787,7 @@ class OracleCompiler(compiler.SQLCompiler):
 
                 # If needed, add the ora_rn, and wrap again with offset.
                 if offset_clause is None:
-                    limitselect._for_update_arg = select._for_update_arg
+                    limitselect._for_update_arg = for_update
                     select = limitselect
                 else:
                     limitselect = limitselect.column(
@@ -786,13 +800,18 @@ class OracleCompiler(compiler.SQLCompiler):
                     offsetselect._oracle_visit = True
                     offsetselect._is_wrapper = True
 
+                    if for_update is not None and for_update.of:
+                        for elem in for_update.of:
+                            if limitselect.corresponding_column(elem) is None:
+                                limitselect.append_column(elem)
+
                     if not self.dialect.use_binds_for_limits:
                         offset_clause = sql.literal_column(
                             "%d" % select._offset)
                     offsetselect.append_whereclause(
                         sql.literal_column("ora_rn") > offset_clause)
 
-                    offsetselect._for_update_arg = select._for_update_arg
+                    offsetselect._for_update_arg = for_update
                     select = offsetselect
 
         return compiler.SQLCompiler.visit_select(self, select, **kwargs)
diff --git a/test/dialect/test_oracle.py b/test/dialect/test_oracle.py
index ed09141bb..1bdddb3bc 100644
--- a/test/dialect/test_oracle.py
+++ b/test/dialect/test_oracle.py
@@ -365,6 +365,79 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
             "mytable_1.myid, mytable_1.name"
         )
 
+    def test_for_update_of_w_limit_adaption_col_present(self):
+        table1 = table('mytable', column('myid'), column('name'))
+
+        self.assert_compile(
+            select([table1.c.myid, table1.c.name]).
+            where(table1.c.myid == 7).
+            with_for_update(nowait=True, of=table1.c.name).
+            limit(10),
+            "SELECT myid, name FROM "
+            "(SELECT mytable.myid AS myid, mytable.name AS name "
+            "FROM mytable WHERE mytable.myid = :myid_1) "
+            "WHERE ROWNUM <= :param_1 FOR UPDATE OF name NOWAIT",
+        )
+
+    def test_for_update_of_w_limit_adaption_col_unpresent(self):
+        table1 = table('mytable', column('myid'), column('name'))
+
+        self.assert_compile(
+            select([table1.c.myid]).
+            where(table1.c.myid == 7).
+            with_for_update(nowait=True, of=table1.c.name).
+            limit(10),
+            "SELECT myid FROM "
+            "(SELECT mytable.myid AS myid, mytable.name AS name "
+            "FROM mytable WHERE mytable.myid = :myid_1) "
+            "WHERE ROWNUM <= :param_1 FOR UPDATE OF name NOWAIT",
+        )
+
+    def test_for_update_of_w_limit_offset_adaption_col_present(self):
+        table1 = table('mytable', column('myid'), column('name'))
+
+        self.assert_compile(
+            select([table1.c.myid, table1.c.name]).
+            where(table1.c.myid == 7).
+            with_for_update(nowait=True, of=table1.c.name).
+            limit(10).offset(50),
+            "SELECT myid, name FROM (SELECT myid, name, ROWNUM AS ora_rn "
+            "FROM (SELECT mytable.myid AS myid, mytable.name AS name "
+            "FROM mytable WHERE mytable.myid = :myid_1) "
+            "WHERE ROWNUM <= :param_1 + :param_2) WHERE ora_rn > :param_2 "
+            "FOR UPDATE OF name NOWAIT",
+        )
+
+    def test_for_update_of_w_limit_offset_adaption_col_unpresent(self):
+        table1 = table('mytable', column('myid'), column('name'))
+
+        self.assert_compile(
+            select([table1.c.myid]).
+            where(table1.c.myid == 7).
+            with_for_update(nowait=True, of=table1.c.name).
+            limit(10).offset(50),
+            "SELECT myid FROM (SELECT myid, ROWNUM AS ora_rn, name "
+            "FROM (SELECT mytable.myid AS myid, mytable.name AS name "
+            "FROM mytable WHERE mytable.myid = :myid_1) "
+            "WHERE ROWNUM <= :param_1 + :param_2) WHERE ora_rn > :param_2 "
+            "FOR UPDATE OF name NOWAIT",
+        )
+
+    def test_for_update_of_w_limit_offset_adaption_partial_col_unpresent(self):
+        table1 = table('mytable', column('myid'), column('foo'), column('bar'))
+
+        self.assert_compile(
+            select([table1.c.myid, table1.c.bar]).
+            where(table1.c.myid == 7).
+            with_for_update(nowait=True, of=[table1.c.foo, table1.c.bar]).
+            limit(10).offset(50),
+            "SELECT myid, bar FROM (SELECT myid, bar, ROWNUM AS ora_rn, "
+            "foo FROM (SELECT mytable.myid AS myid, mytable.bar AS bar, "
+            "mytable.foo AS foo FROM mytable WHERE mytable.myid = :myid_1) "
+            "WHERE ROWNUM <= :param_1 + :param_2) WHERE ora_rn > :param_2 "
+            "FOR UPDATE OF foo, bar NOWAIT"
+        )
+
     def test_limit_preserves_typing_information(self):
         class MyType(TypeDecorator):
             impl = Integer