sql: fix index corruption in udfs with fk cascades

This PR fixes an issue in which UDFs making a FK cascade could cause index corruption. It adds inbound foreign key origin tables to the statement tree, which is used to test for tables with multiple mutations. No release note is necessary, since UDF mutation and FK cascade support is new in this release. Epic: CRDB-25388 Informs: cockroachdb#87289 Release note: None
rharding6373 · Sep 21, 2023 · 2f17874 · 2f17874
1 parent afed367
commit 2f17874
Show file tree

Hide file tree

Showing 5 changed files with 121 additions and 7 deletions.
diff --git a/pkg/sql/logictest/testdata/logic_test/udf_fk b/pkg/sql/logictest/testdata/logic_test/udf_fk
@@ -595,3 +595,93 @@ SELECT * FROM selfref;
 2 2
 
 subtest end
+
+
+subtest corruption_check
+
+statement ok
+DROP TABLE IF EXISTS parent CASCADE;
+
+statement ok
+DROP TABLE IF EXISTS child CASCADE;
+
+statement ok
+DROP TABLE IF EXISTS grandchild CASCADE;
+
+statement ok
+CREATE TABLE parent (j INT PRIMARY KEY);
+
+statement ok
+CREATE TABLE child (i INT PRIMARY KEY, j INT REFERENCES parent (j) ON UPDATE CASCADE ON DELETE CASCADE, INDEX (j));
+
+statement ok
+INSERT INTO parent VALUES (0), (2), (4);
+
+statement ok
+INSERT INTO child VALUES (0, 0);
+
+statement ok
+CREATE OR REPLACE FUNCTION f(k INT) RETURNS INT AS $$
+  UPDATE parent SET j = j + 1 WHERE j = k RETURNING j
+$$ LANGUAGE SQL;
+
+# Check 1 level of cascades.
+statement error pgcode 0A000 pq: multiple mutations of the same table "child" are not supported unless they all use INSERT without ON CONFLICT; this is to prevent data corruption, see documentation of sql.multiple_modifications_of_table.enabled
+WITH x AS (SELECT f(0) AS j), y AS (UPDATE child SET j = 2 WHERE i = 0 RETURNING j) SELECT * FROM x;
+
+query II rowsort
+SELECT i, j FROM child@primary;
+----
+0  0
+
+query II rowsort
+SELECT i, j FROM child@child_j_idx;
+----
+0  0
+
+statement ok
+DROP TABLE IF EXISTS child CASCADE;
+
+statement ok
+CREATE TABLE child (i INT PRIMARY KEY, j INT UNIQUE REFERENCES parent (j) ON UPDATE CASCADE ON DELETE CASCADE, INDEX (j));
+
+statement ok
+CREATE TABLE grandchild (i INT PRIMARY KEY, j INT REFERENCES child (j) ON UPDATE CASCADE ON DELETE CASCADE, INDEX (j));
+
+statement ok
+INSERT INTO child VALUES (0, 0);
+
+statement ok
+INSERT INTO grandchild VALUES (0,0)
+
+# Check 2 levels of cascades.
+statement error pgcode 0A000 pq: multiple mutations of the same table "grandchild" are not supported unless they all use INSERT without ON CONFLICT; this is to prevent data corruption, see documentation of sql.multiple_modifications_of_table.enabled
+WITH x AS (SELECT f(0) AS j), y AS (UPDATE grandchild SET j = 2 WHERE i = 0 RETURNING j) SELECT * FROM x;
+
+statement ok
+DROP TABLE IF EXISTS child CASCADE;
+
+statement ok
+DROP TABLE IF EXISTS grandchild CASCADE;
+
+statement ok
+CREATE TABLE child (i INT PRIMARY KEY, j INT UNIQUE REFERENCES parent (j), INDEX (j));
+
+statement ok
+INSERT INTO child VALUES (0,4)
+
+# Check that we can mutate if there are no actions.
+statement ok
+WITH x AS (SELECT f(0) AS j), y AS (UPDATE child SET j = 2 WHERE i = 0 RETURNING j) SELECT * FROM x;
+
+query II rowsort
+SELECT i, j FROM child@primary;
+----
+0  2
+
+query II rowsort
+SELECT i, j FROM child@child_j_idx;
+----
+0  2
+
+subtest end
diff --git a/pkg/sql/opt/optbuilder/delete.go b/pkg/sql/opt/optbuilder/delete.go
@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/util/intsets"
 )
 
 // buildDelete builds a memo group for a DeleteOp expression, which deletes all
@@ -72,7 +73,8 @@ func (b *Builder) buildDelete(del *tree.Delete, inScope *scope) (outScope *scope
 	b.checkPrivilege(depName, tab, privilege.SELECT)
 
 	// Check if this table has already been mutated in another subquery.
-	b.checkMultipleMutations(tab, generalMutation)
+	var visited intsets.Fast
+	b.checkMultipleMutations(tab, generalMutation, visited)
 
 	var mb mutationBuilder
 	mb.init(b, "delete", tab, alias)

diff --git a/pkg/sql/opt/optbuilder/insert.go b/pkg/sql/opt/optbuilder/insert.go
@@ -218,7 +218,8 @@ func (b *Builder) buildInsert(ins *tree.Insert, inScope *scope) (outScope *scope
 	if ins.OnConflict == nil {
 		mutType = simpleInsert
 	}
-	b.checkMultipleMutations(tab, mutType)
+	var visited intsets.Fast
+	b.checkMultipleMutations(tab, mutType, visited)
 
 	var mb mutationBuilder
 	if ins.OnConflict != nil && ins.OnConflict.IsUpsertAlias() {

diff --git a/pkg/sql/opt/optbuilder/update.go b/pkg/sql/opt/optbuilder/update.go
@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/intsets"
 	"github.com/cockroachdb/errors"
 )
 
@@ -86,7 +87,8 @@ func (b *Builder) buildUpdate(upd *tree.Update, inScope *scope) (outScope *scope
 	b.checkPrivilege(depName, tab, privilege.SELECT)
 
 	// Check if this table has already been mutated in another subquery.
-	b.checkMultipleMutations(tab, generalMutation)
+	var visited intsets.Fast
+	b.checkMultipleMutations(tab, generalMutation, visited)
 
 	var mb mutationBuilder
 	mb.init(b, "update", tab, alias)

diff --git a/pkg/sql/opt/optbuilder/util.go b/pkg/sql/opt/optbuilder/util.go
@@ -11,6 +11,7 @@
 package optbuilder
 
 import (
+	"context"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
@@ -23,6 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/intsets"
 	"github.com/cockroachdb/errors"
 )
 
@@ -512,17 +514,34 @@ func (b *Builder) resolveSchemaForCreate(
 	return sch, resName
 }
 
-func (b *Builder) checkMultipleMutations(tab cat.Table, typ mutationType) {
-	if !b.stmtTree.CanMutateTable(tab.ID(), typ) &&
-		!multipleModificationsOfTableEnabled.Get(&b.evalCtx.Settings.SV) &&
-		!b.evalCtx.SessionData().MultipleModificationsOfTable {
+func (b *Builder) checkMultipleMutations(tab cat.Table, typ mutationType, visited intsets.Fast) {
+	if multipleModificationsOfTableEnabled.Get(&b.evalCtx.Settings.SV) ||
+		b.evalCtx.SessionData().MultipleModificationsOfTable {
+		return
+	}
+	if visited.Contains(int(tab.ID())) {
+		return
+	}
+	if !b.stmtTree.CanMutateTable(tab.ID(), typ) {
 		panic(pgerror.Newf(
 			pgcode.FeatureNotSupported,
 			"multiple mutations of the same table %q are not supported unless they all "+
 				"use INSERT without ON CONFLICT; this is to prevent data corruption, see "+
 				"documentation of sql.multiple_modifications_of_table.enabled", tab.Name(),
 		))
 	}
+	visited.Add(int(tab.ID()))
+
+	// If this table references foreign keys that will also be mutated, then add
+	// them to the statement tree via a recursive call.
+	for i := 0; i < tab.InboundForeignKeyCount(); i++ {
+		fk := tab.InboundForeignKey(i)
+		if fk.DeleteReferenceAction() != tree.NoAction ||
+			fk.UpdateReferenceAction() != tree.NoAction {
+			fkTab := resolveTable(context.Background(), b.catalog, fk.OriginTableID())
+			b.checkMultipleMutations(fkTab, typ, visited)
+		}
+	}
 }
 
 // resolveTableForMutation is a helper method for building mutations. It returns