Repository: edgedb/edgedb

Issue:
Properly reject correlation with FOR iterator containing DML
We should reject queries like `SELECT (Person, (FOR x in {Person} UNION ((INSERT Note {name := x.name}))));`

Currently we generate invalid SQL.


Discussion:
https://github.com/edgedb/edgedb/pull/1757/ made this succeed by losing the correlation between the two Persons, which is also a bug!

Resolving patch:
diff --git a/edb/edgeql/compiler/stmt.py b/edb/edgeql/compiler/stmt.py
--- a/edb/edgeql/compiler/stmt.py
+++ b/edb/edgeql/compiler/stmt.py
@@ -44,6 +44,7 @@
 from edb.schema import types as s_types
 
 from edb.edgeql import ast as qlast
+from edb.edgeql import utils as qlutils
 from edb.edgeql import qltypes
 
 from . import astutils
@@ -167,7 +168,19 @@ def compile_ForQuery(
         sctx.iterator_path_ids |= {stmt.iterator_stmt.path_id}
         node = ictx.path_scope.find_descendant(iterator_stmt.path_id)
         if node is not None:
-            node.attach_subtree(view_scope_info.path_scope)
+            # If the body contains DML, then we need to prohibit
+            # correlation between the iterator and the enclosing
+            # query, since the correlation imposes compilation issues
+            # we aren't willing to tackle.
+            # Do this by sticking the iterator subtree onto a branch
+            # with a factoring fence.
+            if qlutils.contains_dml(qlstmt.result):
+                node = node.attach_branch()
+                node.factoring_fence = True
+                node = node.attach_branch()
+
+            node.attach_subtree(view_scope_info.path_scope,
+                                context=iterator.context)
 
         # Compile the body
         stmt.result = compile_result_clause(
diff --git a/edb/edgeql/utils.py b/edb/edgeql/utils.py
--- a/edb/edgeql/utils.py
+++ b/edb/edgeql/utils.py
@@ -121,3 +121,17 @@ def inline_anchors(
 
     inliner = AnchorInliner(anchors)
     inliner.visit(ql_expr)
+
+
+def contains_dml(ql_expr: qlast.Base) -> bool:
+    """Check whether a expression contains any DML in a subtree."""
+    # If this ends up being a perf problem, we can use a visitor
+    # directly and cache.
+    dml_types = (qlast.InsertQuery, qlast.UpdateQuery, qlast.DeleteQuery)
+    if isinstance(ql_expr, dml_types):
+        return True
+
+    res = ast.find_children(ql_expr, lambda x: isinstance(x, dml_types),
+                            terminate_early=True)
+
+    return bool(res)


Test patch:
diff --git a/tests/test_edgeql_insert.py b/tests/test_edgeql_insert.py
--- a/tests/test_edgeql_insert.py
+++ b/tests/test_edgeql_insert.py
@@ -1188,6 +1188,57 @@ async def test_edgeql_insert_for_17(self):
             ],
         )
 
+    async def test_edgeql_insert_for_bad_01(self):
+        with self.assertRaisesRegex(
+            edgedb.errors.QueryError,
+            "cannot reference correlated set",
+        ):
+            await self.con.execute("""
+                WITH MODULE test
+                SELECT (Person,
+                        (FOR x in {Person} UNION (
+                             INSERT Note {name := x.name})));
+            """)
+
+    async def test_edgeql_insert_for_bad_02(self):
+        with self.assertRaisesRegex(
+            edgedb.errors.QueryError,
+            "cannot reference correlated set",
+        ):
+            await self.con.execute("""
+                WITH MODULE test
+                SELECT (Person,
+                        (FOR x in {Person} UNION (
+                             SELECT (INSERT Note {name := x.name}))));
+            """)
+
+    async def test_edgeql_insert_for_bad_03(self):
+        with self.assertRaisesRegex(
+            edgedb.errors.QueryError,
+            "cannot reference correlated set",
+        ):
+            await self.con.execute("""
+                WITH MODULE test
+                SELECT ((FOR x in {Person} UNION (
+                             INSERT Note {name := x.name})),
+                        Person);
+            """)
+
+    async def test_edgeql_insert_for_bad_04(self):
+        with self.assertRaisesRegex(
+            edgedb.errors.QueryError,
+            "cannot reference correlated set",
+        ):
+            await self.con.execute("""
+                WITH MODULE test
+                SELECT (Person,
+                        (FOR x in {Person} UNION (
+                             SELECT (
+                                 20,
+                                 (FOR y in {"hello", "world"} UNION (
+                                  INSERT Note {name := y ++ x.name}))))));
+            """)
+
     async def test_edgeql_insert_default_01(self):
         await self.con.execute(r'''
             # create 10 DefaultTest3 objects, each object is defined