sql: support foreign key checks in udfs

Before this change, we would return an error if a UDF attempted to make
a foreign key check because checks were not supported in routines.

This change adds support for running postquery checks, like FK checks,
in routines. It leverages the existing DistSQL postquery planner to run
checks built for a routine's statement after the statement has been
planned and run. This also allows UDFs to take advantage of the parallel
FK check capabilities. Note that foreign key checks are run after each
statement in the UDF body, as well as after the main query if required.

This change also refactors how extendedEvalContexts are copied for
parallel checks.

Support for FK cascades will come in a later PR.

Epic: CRDB-25388
Informs: #87289

Release note: None

pkg/sql/apply_join.go

@@ -12,6 +12,7 @@ package sql
 
 import (
 	"context"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"strconv"
 	"sync/atomic"
 
@@ -301,12 +302,16 @@ func runPlanInsidePlan(
 		// Create a separate memory account for the results of the subqueries.
 		// Note that we intentionally defer the closure of the account until we
 		// return from this method (after the main query is executed).
+		subqueryEvalCtxFactory := func() *extendedEvalContext {
+			return params.p.ExtendedEvalContextCopyAndReset()
+		}
+
 		subqueryResultMemAcc := params.p.Mon().MakeBoundAccount()
 		defer subqueryResultMemAcc.Close(ctx)
 		if !execCfg.DistSQLPlanner.PlanAndRunSubqueries(
 			ctx,
 			params.p,
-			params.extendedEvalCtx.copy,
+			subqueryEvalCtxFactory,
 			plan.subqueryPlans,
 			recv,
 			&subqueryResultMemAcc,
@@ -341,9 +346,51 @@ func runPlanInsidePlan(
 
 	finishedSetupFn, cleanup := getFinishedSetupFn(&plannerCopy)
 	defer cleanup()
+	var evalCtxFactory func(usedConcurrently bool) *extendedEvalContext
+	if len(plan.cascades) != 0 ||
+		len(plan.checkPlans) != 0 {
+		serialEvalCtx := plannerCopy.ExtendedEvalContextCopyAndReset()
+		evalCtxFactory = func(usedConcurrently bool) *extendedEvalContext {
+			if usedConcurrently {
+				return plannerCopy.ExtendedEvalContextCopyAndReset()
+			}
+			// Reuse the same object if this factory is not used concurrently.
+			plannerCopy.ExtendedEvalContextReset(serialEvalCtx)
+			return serialEvalCtx
+		}
+	}
 	execCfg.DistSQLPlanner.PlanAndRun(
 		ctx, evalCtx, planCtx, plannerCopy.Txn(), plan.main, recv, finishedSetupFn,
 	)
+	if p := planCtx.getPortalPauseInfo(); p != nil {
+		if buildutil.CrdbTestBuild && p.resumableFlow.flow == nil {
+			checkErr := errors.AssertionFailedf("flow for portal %s cannot be found", plannerCopy.pausablePortal.Name)
+			if recv.commErr != nil {
+				recv.commErr = errors.CombineErrors(recv.commErr, checkErr)
+			} else {
+				return checkErr
+			}
+		}
+		if !p.resumableFlow.cleanup.isComplete {
+			p.resumableFlow.cleanup.appendFunc(namedFunc{
+				fName: "cleanup flow", f: func() {
+					p.resumableFlow.flow.Cleanup(ctx)
+				},
+			})
+		}
+	}
+
+	if recv.commErr != nil || recv.getError() != nil {
+		return recv.commErr
+	}
+
+	execCfg.DistSQLPlanner.PlanAndRunCascadesAndChecks(
+		ctx, &plannerCopy, evalCtxFactory, &plannerCopy.curPlan.planComponents, recv,
+	)
+	if recv.commErr != nil {
+		return recv.commErr
+	}
+
 	return resultWriter.Err()
 }
 

pkg/sql/conn_executor_exec.go

@@ -2078,19 +2078,17 @@ func (ex *connExecutor) execWithDistSQLEngine(
 		if len(planner.curPlan.subqueryPlans) != 0 ||
 			len(planner.curPlan.cascades) != 0 ||
 			len(planner.curPlan.checkPlans) != 0 {
-			var serialEvalCtx extendedEvalContext
-			ex.initEvalCtx(ctx, &serialEvalCtx, planner)
+			serialEvalCtx := planner.ExtendedEvalContextCopyAndReset()
+			ex.initEvalCtx(ctx, serialEvalCtx, planner)
 			evalCtxFactory = func(usedConcurrently bool) *extendedEvalContext {
 				// Reuse the same object if this factory is not used concurrently.
-				factoryEvalCtx := &serialEvalCtx
+				factoryEvalCtx := serialEvalCtx
 				if usedConcurrently {
-					factoryEvalCtx = &extendedEvalContext{}
+					factoryEvalCtx = planner.ExtendedEvalContextCopyAndReset()
 					ex.initEvalCtx(ctx, factoryEvalCtx, planner)
 				}
 				ex.resetEvalCtx(factoryEvalCtx, planner.txn, planner.ExtendedEvalContext().StmtTimestamp)
-				factoryEvalCtx.Placeholders = &planner.semaCtx.Placeholders
-				factoryEvalCtx.Annotations = &planner.semaCtx.Annotations
-				factoryEvalCtx.SessionID = planner.ExtendedEvalContext().SessionID
+				planner.ExtendedEvalContextReset(factoryEvalCtx)
 				return factoryEvalCtx
 			}
 		}

pkg/sql/logictest/testdata/logic_test/udf_fk

@@ -0,0 +1,236 @@
+# Disable fast path for some test runs.
+let $enable_insert_fast_path
+SELECT random() < 0.5
+
+statement ok
+SET enable_insert_fast_path = $enable_insert_fast_path
+
+statement ok
+CREATE TABLE parent (p INT PRIMARY KEY);
+
+statement ok
+CREATE TABLE child (c INT PRIMARY KEY, p INT NOT NULL REFERENCES parent(p));
+
+
+subtest insert
+
+statement ok
+CREATE FUNCTION f_fk_c(k INT, r INT) RETURNS RECORD AS $$
+  INSERT INTO child VALUES (k,r) RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_fk_p(r INT) RETURNS RECORD AS $$
+  INSERT INTO parent VALUES (r) RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_fk_c_p(k INT, r INT) RETURNS RECORD AS $$
+  INSERT INTO child VALUES (k,r);
+  INSERT INTO parent VALUES (r) RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_fk_p_c(k INT, r INT) RETURNS RECORD AS $$
+  INSERT INTO parent VALUES (r);
+  INSERT INTO child VALUES (k, r) RETURNING *;
+$$ LANGUAGE SQL;
+
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c(100, 1);
+
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_p(100, 1);
+
+query T
+SELECT f_fk_p_c(100, 1);
+----
+(100,1)
+
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+WITH x AS (SELECT f_fk_c(101, 2)) INSERT INTO parent VALUES (2);
+
+query T
+WITH x AS (INSERT INTO parent VALUES (2) RETURNING p) SELECT f_fk_c(101, 2);
+----
+(101,2)
+
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c(102, 3), f_fk_p(3);
+
+query TT
+SELECT f_fk_p(3), f_fk_c(102, 3);
+----
+(3) (102,3)
+
+statement ok
+TRUNCATE parent CASCADE
+
+statement ok
+INSERT INTO parent (p) VALUES (1);
+
+statement ok
+CREATE FUNCTION f_fk_c_multi(k1 INT, r1 INT, k2 INT, r2 INT) RETURNS SETOF RECORD AS $$
+  INSERT INTO child VALUES (k1,r1);
+  INSERT INTO child VALUES (k2,r2);
+  SELECT * FROM child WHERE c = k1 OR c = k2;
+$$ LANGUAGE SQL;
+
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_multi(101, 1, 102, 2);
+
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_multi(101, 2, 102, 1);
+
+query T rowsort
+SELECT f_fk_c_multi(101, 1, 102, 1);
+----
+(101,1)
+(102,1)
+
+# Sequences advance even if subsequent statements fail foreign key checks.
+statement ok
+CREATE SEQUENCE s;
+
+statement ok
+CREATE FUNCTION f_fk_c_seq_first(k INT, r INT) RETURNS RECORD AS $$
+  SELECT nextval('s');
+  INSERT INTO child VALUES (k,r) RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_fk_c_seq_last(k INT, r INT) RETURNS RECORD AS $$
+  INSERT INTO child VALUES (k,r) RETURNING *;
+  SELECT nextval('s');
+$$ LANGUAGE SQL;
+
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_seq_last(103,2);
+
+statement error pq: currval\(\): currval of sequence \"test.public.s\" is not yet defined in this session
+SELECT currval('s');
+
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_seq_first(103,2);
+
+query I
+SELECT currval('s');
+----
+1
+
+subtest delete
+
+statement ok
+TRUNCATE parent CASCADE
+
+statement ok
+INSERT INTO parent (p) VALUES (1), (2), (3), (4);
+
+statement ok
+INSERT INTO child (c, p) VALUES (100, 1), (101, 2), (102, 3);
+
+query I rowsort
+SELECT * FROM parent
+----
+1
+2
+3
+4
+
+query II rowsort
+SELECT * FROM child
+----
+100 1
+101 2
+102 3
+
+statement ok
+CREATE FUNCTION f_fk_c_del(k INT) RETURNS RECORD AS $$
+  DELETE FROM child WHERE c = k RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_fk_p_del(r INT) RETURNS RECORD AS $$
+  DELETE FROM parent WHERE p = r RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_fk_c_p_del(k INT, r INT) RETURNS RECORD AS $$
+  DELETE FROM child WHERE c = k RETURNING *;
+  DELETE FROM parent WHERE p = r RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_fk_p_c_del(k INT, r INT) RETURNS RECORD AS $$
+  DELETE FROM parent WHERE p = r RETURNING *;
+  DELETE FROM child WHERE c = k RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+SELECT f_fk_p_del(4);
+
+statement error pq: delete on table "parent" violates foreign key constraint "child_p_fkey" on table "child"\nDETAIL: Key \(p\)=\(3\) is still referenced from table "child"\.
+SELECT f_fk_p_del(3);
+
+statement ok
+SELECT f_fk_c_del(102);
+
+statement ok
+SELECT f_fk_p_del(3);
+
+statement error pq: delete on table "parent" violates foreign key constraint "child_p_fkey" on table "child"\nDETAIL: Key \(p\)=\(2\) is still referenced from table "child"\.
+SELECT f_fk_p_c_del(101,2);
+
+statement ok
+SELECT f_fk_c_p_del(101,2);
+
+statement ok
+SELECT f_fk_c_del(100), f_fk_p_del(1);
+
+query I rowsort
+SELECT * FROM parent
+----
+
+query II rowsort
+SELECT * FROM child
+----
+
+
+subtest upsert
+
+statement ok
+TRUNCATE parent CASCADE
+
+statement ok
+CREATE FUNCTION f_fk_c_ocdu(k INT, r INT) RETURNS RECORD AS $$
+  INSERT INTO child VALUES (k, r) ON CONFLICT (c) DO UPDATE SET p = r RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+INSERT INTO parent VALUES (1), (3);
+
+# Insert
+statement ok
+SELECT f_fk_c_ocdu(100,1);
+
+# Update to value not in parent fails.
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_ocdu(100,2);
+
+# Inserting value not in parent fails.
+statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_ocdu(101,2);
+
+statement ok
+CREATE FUNCTION f_fk_c_ups(k INT, r INT) RETURNS RECORD AS $$
+  UPSERT INTO child VALUES (k, r) RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+SELECT f_fk_c_ups(102,3);
+
+statement error pq: upsert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_ups(102,4);
+
+statement error pq: upsert on table "child" violates foreign key constraint "child_p_fkey"
+SELECT f_fk_c_ups(103,4);