opt: disallow mutations under union

Because the inputs of a UNION or UNION ALL are run in parallel, it is
not possible for either of them to be mutations. The right way to do
this is to use WITH (above the union). The optimizer will have code to
pull up mutations into top-level WITHs, but until then we want to
disallow such queries from executing. This change adds a check for
this condition in the execbuilder.

Informs cockroachdb#40853.

Release note (sql change): Mutations under UNION or UNION ALL are
now disallowed; WITH should be used on top of the union operation
instead. This restriction is temporary and will be lifted in a future
release.

Release justification: low-risk fix for high-severity bug (category
3).

pkg/sql/logictest/testdata/logic_test/union

@@ -286,3 +286,28 @@ SELECT a FROM (SELECT a FROM c UNION ALL SELECT a FROM a) WHERE a > 0 AND a < 3
 ----
 1
 1
+
+# Ensure that mutations under a UNION or UNION ALL error out (#40853).
+statement error mutations not supported under UNION
+SELECT * FROM uniontest
+UNION SELECT * FROM [INSERT INTO uniontest VALUES (1), (1) RETURNING k, v]
+
+statement error mutations not supported under UNION
+SELECT * FROM [INSERT INTO uniontest VALUES (1), (1) RETURNING k, v]
+UNION ALL SELECT * FROM uniontest
+
+statement error mutations not supported under UNION
+SELECT * FROM uniontest
+UNION SELECT * FROM (
+  WITH cte AS (INSERT INTO uniontest VALUES (1), (1) RETURNING k, v) SELECT * FROM cte
+)
+
+# The right way to run such a query is to use WITH above the union.
+statement ok
+WITH cte AS (INSERT INTO uniontest VALUES (1), (1) RETURNING k, v)
+(SELECT * FROM cte UNION SELECT * FROM uniontest)
+
+# Other set ops are allowed.
+statement ok
+SELECT * FROM uniontest
+EXCEPT SELECT * FROM [INSERT INTO uniontest VALUES (1), (1) RETURNING k, v]

pkg/sql/logictest/testdata/logic_test/upsert

@@ -373,12 +373,13 @@ upsert INTO upsert_returning VALUES (1, 5, 4, 3), (6, 5, 4, 3) RETURNING *
 statement ok
 COMMIT
 
-# For #22300. Test UPSERT ... RETURNING with UNION.
-query I rowsort
-SELECT a FROM [UPSERT INTO upsert_returning VALUES (7) RETURNING a] UNION VALUES (8)
-----
-7
-8
+# TODO(justin): reenable when we fix #35060 / #40853.
+## For #22300. Test UPSERT ... RETURNING with UNION.
+#query I rowsort
+#SELECT a FROM [UPSERT INTO upsert_returning VALUES (7) RETURNING a] UNION VALUES (8)
+#----
+#7
+#8
 
 # For #6710. Add an unused column to disable the fast path which doesn't have this bug.
 statement ok

pkg/sql/opt/exec/execbuilder/relational.go

@@ -1110,6 +1110,13 @@ func (b *Builder) buildSetOp(set memo.RelExpr) (execPlan, error) {
 	default:
 		panic(errors.AssertionFailedf("invalid operator %s", log.Safe(set.Op())))
 	}
+	// TODO(justin): We cannot execute mutations under a UNION or UNION ALL
+	// (because mutations can't run in parallel). Once we pull up all mutations
+	// into top-level With expressions, this case will not be possible anymore.
+	if typ == tree.UnionOp && (leftExpr.Relational().CanMutate || rightExpr.Relational().CanMutate) {
+		err := unimplemented.NewWithIssuef(40853, "mutations not supported under UNION; use WITH instead")
+		return execPlan{}, err
+	}
 
 	node, err := b.factory.ConstructSetOp(typ, all, left.root, right.root)
 	if err != nil {