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

1 files changed, 57 insertions, 1 deletions

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