Include UPDATE/DELETE extra_froms in correlation

Fixed bug where the multi-table support for UPDATE and DELETE statements
did not consider the additional FROM elements as targets for correlation,
when a correlated SELECT were also combined with the statement.  This
change now includes that a SELECT statement in the WHERE clause for such a
statement will try to auto-correlate back to these additional tables in the
parent UPDATE/DELETE or unconditionally correlate if
:meth:`.Select.correlate` is used.  Note that auto-correlation raises an
error if the SELECT statement would have no FROM clauses as a result, which
can now occur if the parent UPDATE/DELETE specifies the same tables in its
additional set of tables ; specify :meth:`.Select.correlate` explicitly to
resolve.

Change-Id: Ie11eaad7e49af3f59df11691b104d6359341bdae
Fixes: #4313

2 files changed, 131 insertions, 3 deletions

diff --git a/test/sql/test_delete.py b/test/sql/test_delete.py
index 7d18db9c9..91f2c2cdc 100644
--- a/test/sql/test_delete.py
+++ b/test/sql/test_delete.py
@@ -1,10 +1,13 @@
 #! coding:utf-8
 
 from sqlalchemy import Integer, String, ForeignKey, delete, select, and_, \
-    or_
+    or_, exists
 from sqlalchemy.dialects import mysql
+from sqlalchemy.engine import default
 from sqlalchemy import testing
-from sqlalchemy.testing import AssertsCompiledSQL, fixtures, eq_
+from sqlalchemy import exc
+from sqlalchemy.testing import AssertsCompiledSQL, fixtures, eq_, \
+    assert_raises_message
 from sqlalchemy.testing.schema import Table, Column
 
 
@@ -103,6 +106,75 @@ class DeleteTest(_DeleteTestBase, fixtures.TablesTest, AssertsCompiledSQL):
                             ')')
 
 
+class DeleteFromCompileTest(
+        _DeleteTestBase, fixtures.TablesTest, AssertsCompiledSQL):
+    # DELETE FROM is also tested by individual dialects since there is no
+    # consistent syntax.  here we use the StrSQLcompiler which has a fake
+    # syntax.
+
+    __dialect__ = 'default_enhanced'
+
+    def test_delete_extra_froms(self):
+        table1, table2 = self.tables.mytable, self.tables.myothertable
+
+        stmt = table1.delete().where(table1.c.myid == table2.c.otherid)
+        self.assert_compile(
+            stmt,
+            "DELETE FROM mytable , myothertable "
+            "WHERE mytable.myid = myothertable.otherid",
+        )
+
+    def test_correlation_to_extra(self):
+        table1, table2 = self.tables.mytable, self.tables.myothertable
+
+        stmt = table1.delete().where(
+            table1.c.myid == table2.c.otherid).where(
+                ~exists().where(table2.c.otherid == table1.c.myid).
+                where(table2.c.othername == 'x').correlate(table2)
+        )
+
+        self.assert_compile(
+            stmt,
+            "DELETE FROM mytable , myothertable WHERE mytable.myid = "
+            "myothertable.otherid AND NOT (EXISTS "
+            "(SELECT * FROM mytable WHERE myothertable.otherid = "
+            "mytable.myid AND myothertable.othername = :othername_1))",
+        )
+
+    def test_dont_correlate_to_extra(self):
+        table1, table2 = self.tables.mytable, self.tables.myothertable
+
+        stmt = table1.delete().where(
+            table1.c.myid == table2.c.otherid).where(
+                ~exists().where(table2.c.otherid == table1.c.myid).
+                where(table2.c.othername == 'x').correlate()
+        )
+
+        self.assert_compile(
+            stmt,
+            "DELETE FROM mytable , myothertable WHERE mytable.myid = "
+            "myothertable.otherid AND NOT (EXISTS "
+            "(SELECT * FROM myothertable, mytable "
+            "WHERE myothertable.otherid = "
+            "mytable.myid AND myothertable.othername = :othername_1))",
+        )
+
+    def test_autocorrelate_error(self):
+        table1, table2 = self.tables.mytable, self.tables.myothertable
+
+        stmt = table1.delete().where(
+            table1.c.myid == table2.c.otherid).where(
+                ~exists().where(table2.c.otherid == table1.c.myid).
+                where(table2.c.othername == 'x')
+        )
+
+        assert_raises_message(
+            exc.InvalidRequestError,
+            ".*returned no FROM clauses due to auto-correlation.*",
+            stmt.compile, dialect=default.StrCompileDialect()
+        )
+
+
 class DeleteFromRoundTripTest(fixtures.TablesTest):
     __backend__ = True
 
diff --git a/test/sql/test_update.py b/test/sql/test_update.py
index cc5b4962b..138581061 100644
--- a/test/sql/test_update.py
+++ b/test/sql/test_update.py
@@ -1,5 +1,5 @@
 from sqlalchemy import Integer, String, ForeignKey, and_, or_, func, \
-    literal, update, table, bindparam, column, select, exc
+    literal, update, table, bindparam, column, select, exc, exists
 from sqlalchemy import testing
 from sqlalchemy.dialects import mysql
 from sqlalchemy.engine import default
@@ -591,6 +591,62 @@ class UpdateFromCompileTest(_UpdateFromTestBase, fixtures.TablesTest,
             'AND anon_1.email_address = :email_address_1',
             checkparams=checkparams)
 
+    def test_correlation_to_extra(self):
+        users, addresses = self.tables.users, self.tables.addresses
+
+        stmt = users.update().values(name="newname").where(
+            users.c.id == addresses.c.user_id
+        ).where(
+            ~exists().where(
+                addresses.c.user_id == users.c.id
+            ).where(addresses.c.email_address == 'foo').correlate(addresses)
+        )
+
+        self.assert_compile(
+            stmt,
+            "UPDATE users SET name=:name FROM addresses WHERE "
+            "users.id = addresses.user_id AND NOT "
+            "(EXISTS (SELECT * FROM users WHERE addresses.user_id = users.id "
+            "AND addresses.email_address = :email_address_1))"
+        )
+
+    def test_dont_correlate_to_extra(self):
+        users, addresses = self.tables.users, self.tables.addresses
+
+        stmt = users.update().values(name="newname").where(
+            users.c.id == addresses.c.user_id
+        ).where(
+            ~exists().where(
+                addresses.c.user_id == users.c.id
+            ).where(addresses.c.email_address == 'foo').correlate()
+        )
+
+        self.assert_compile(
+            stmt,
+            "UPDATE users SET name=:name FROM addresses WHERE "
+            "users.id = addresses.user_id AND NOT "
+            "(EXISTS (SELECT * FROM addresses, users "
+            "WHERE addresses.user_id = users.id "
+            "AND addresses.email_address = :email_address_1))"
+        )
+
+    def test_autocorrelate_error(self):
+        users, addresses = self.tables.users, self.tables.addresses
+
+        stmt = users.update().values(name="newname").where(
+            users.c.id == addresses.c.user_id
+        ).where(
+            ~exists().where(
+                addresses.c.user_id == users.c.id
+            ).where(addresses.c.email_address == 'foo')
+        )
+
+        assert_raises_message(
+            exc.InvalidRequestError,
+            ".*returned no FROM clauses due to auto-correlation.*",
+            stmt.compile, dialect=default.StrCompileDialect()
+        )
+
 
 class UpdateFromRoundTripTest(_UpdateFromTestBase, fixtures.TablesTest):
     __backend__ = True