Skip to content

Commit 54c3e26

Browse files
rharding6373rharding6373
committed
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: cockroachdb#87289 Release note: None
1 parent d6665c0 commit 54c3e26

File tree

14 files changed

+429
-10
lines changed

14 files changed

+429
-10
lines changed

‎pkg/ccl/logictestccl/tests/3node-tenant/generated_test.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/sql/apply_join.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,23 @@ func runPlanInsidePlan(
342342
execCfg.DistSQLPlanner.PlanAndRun(
343343
ctx, evalCtx, planCtx, plannerCopy.Txn(), plan.main, recv, finishedSetupFn,
344344
)
345+
346+
// Check if there was an error interacting with the resultWriter.
347+
if recv.commErr != nil {
348+
return recv.commErr
349+
}
350+
351+
evalCtxFactory2 := func(usedConcurrently bool) *extendedEvalContext {
352+
return evalCtxFactory()
353+
}
354+
355+
execCfg.DistSQLPlanner.PlanAndRunCascadesAndChecks(
356+
ctx, &plannerCopy, evalCtxFactory2, &plannerCopy.curPlan.planComponents, recv,
357+
)
358+
if recv.commErr != nil {
359+
return recv.commErr
360+
}
361+
345362
return resultWriter.Err()
346363
}
347364

‎pkg/sql/conn_executor_exec.go‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2078,19 +2078,17 @@ func (ex *connExecutor) execWithDistSQLEngine(
20782078
if len(planner.curPlan.subqueryPlans) != 0 ||
20792079
len(planner.curPlan.cascades) != 0 ||
20802080
len(planner.curPlan.checkPlans) != 0 {
2081-
var serialEvalCtx extendedEvalContext
2082-
ex.initEvalCtx(ctx, &serialEvalCtx, planner)
2081+
serialEvalCtx := planner.ExtendedEvalContextCopyAndReset()
2082+
ex.initEvalCtx(ctx, serialEvalCtx, planner)
20832083
evalCtxFactory = func(usedConcurrently bool) *extendedEvalContext {
20842084
// Reuse the same object if this factory is not used concurrently.
2085-
factoryEvalCtx := &serialEvalCtx
2085+
factoryEvalCtx := serialEvalCtx
20862086
if usedConcurrently {
2087-
factoryEvalCtx = &extendedEvalContext{}
2087+
factoryEvalCtx = planner.ExtendedEvalContextCopyAndReset()
20882088
ex.initEvalCtx(ctx, factoryEvalCtx, planner)
20892089
}
20902090
ex.resetEvalCtx(factoryEvalCtx, planner.txn, planner.ExtendedEvalContext().StmtTimestamp)
2091-
factoryEvalCtx.Placeholders = &planner.semaCtx.Placeholders
2092-
factoryEvalCtx.Annotations = &planner.semaCtx.Annotations
2093-
factoryEvalCtx.SessionID = planner.ExtendedEvalContext().SessionID
2091+
planner.ExtendedEvalContextReset(factoryEvalCtx)
20942092
return factoryEvalCtx
20952093
}
20962094
}

‎pkg/sql/logictest/testdata/logic_test/udf_fk‎

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,228 @@
1+
# Disable fast path for some test runs.
2+
let $enable_insert_fast_path
3+
SELECT random() < 0.5
4+
5+
statement ok
6+
SET enable_insert_fast_path = $enable_insert_fast_path
7+
8+
statement ok
9+
CREATE TABLE parent (p INT PRIMARY KEY);
10+
11+
statement ok
12+
CREATE TABLE child (c INT PRIMARY KEY, p INT NOT NULL REFERENCES parent(p));
13+
14+
15+
subtest insert
16+
17+
statement ok
18+
CREATE FUNCTION f_fk_c(k INT, r INT) RETURNS RECORD AS $$
19+
INSERT INTO child VALUES (k,r) RETURNING *;
20+
$$ LANGUAGE SQL;
21+
22+
statement ok
23+
CREATE FUNCTION f_fk_p(r INT) RETURNS RECORD AS $$
24+
INSERT INTO parent VALUES (r) RETURNING *;
25+
$$ LANGUAGE SQL;
26+
27+
statement ok
28+
CREATE FUNCTION f_fk_c_p(k INT, r INT) RETURNS RECORD AS $$
29+
INSERT INTO child VALUES (k,r);
30+
INSERT INTO parent VALUES (r) RETURNING *;
31+
$$ LANGUAGE SQL;
32+
33+
statement ok
34+
CREATE FUNCTION f_fk_p_c(k INT, r INT) RETURNS RECORD AS $$
35+
INSERT INTO parent VALUES (r);
36+
INSERT INTO child VALUES (k, r) RETURNING *;
37+
$$ LANGUAGE SQL;
38+
39+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
40+
SELECT f_fk_c(100, 1);
41+
42+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
43+
SELECT f_fk_c_p(100, 1);
44+
45+
query T
46+
SELECT f_fk_p_c(100, 1);
47+
----
48+
(100,1)
49+
50+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
51+
WITH x AS (SELECT f_fk_c(101, 2)) INSERT INTO parent VALUES (2);
52+
53+
query T
54+
WITH x AS (INSERT INTO parent VALUES (2) RETURNING p) SELECT f_fk_c(101, 2);
55+
----
56+
(101,2)
57+
58+
statement ok
59+
TRUNCATE parent CASCADE
60+
61+
statement ok
62+
INSERT INTO parent (p) VALUES (1);
63+
64+
statement ok
65+
CREATE FUNCTION f_fk_c_multi(k1 INT, r1 INT, k2 INT, r2 INT) RETURNS SETOF RECORD AS $$
66+
INSERT INTO child VALUES (k1,r1);
67+
INSERT INTO child VALUES (k2,r2);
68+
SELECT * FROM child WHERE c = k1 OR c = k2;
69+
$$ LANGUAGE SQL;
70+
71+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
72+
SELECT f_fk_c_multi(101, 1, 102, 2);
73+
74+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
75+
SELECT f_fk_c_multi(101, 2, 102, 1);
76+
77+
query T rowsort
78+
SELECT f_fk_c_multi(101, 1, 102, 1);
79+
----
80+
(101,1)
81+
(102,1)
82+
83+
# Sequences advance even if subsequent statements fail foreign key checks.
84+
statement ok
85+
CREATE SEQUENCE s;
86+
87+
statement ok
88+
CREATE FUNCTION f_fk_c_seq_first(k INT, r INT) RETURNS RECORD AS $$
89+
SELECT nextval('s');
90+
INSERT INTO child VALUES (k,r) RETURNING *;
91+
$$ LANGUAGE SQL;
92+
93+
statement ok
94+
CREATE FUNCTION f_fk_c_seq_last(k INT, r INT) RETURNS RECORD AS $$
95+
INSERT INTO child VALUES (k,r) RETURNING *;
96+
SELECT nextval('s');
97+
$$ LANGUAGE SQL;
98+
99+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
100+
SELECT f_fk_c_seq_last(103,2);
101+
102+
statement error pq: currval\(\): currval of sequence \"test.public.s\" is not yet defined in this session
103+
SELECT currval('s');
104+
105+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
106+
SELECT f_fk_c_seq_first(103,2);
107+
108+
query I
109+
SELECT currval('s');
110+
----
111+
1
112+
113+
subtest delete
114+
115+
statement ok
116+
TRUNCATE parent CASCADE
117+
118+
statement ok
119+
INSERT INTO parent (p) VALUES (1), (2), (3), (4);
120+
121+
statement ok
122+
INSERT INTO child (c, p) VALUES (100, 1), (101, 2), (102, 3);
123+
124+
query I rowsort
125+
SELECT * FROM parent
126+
----
127+
1
128+
2
129+
3
130+
4
131+
132+
query II rowsort
133+
SELECT * FROM child
134+
----
135+
100 1
136+
101 2
137+
102 3
138+
139+
statement ok
140+
CREATE FUNCTION f_fk_c_del(k INT) RETURNS RECORD AS $$
141+
DELETE FROM child WHERE c = k RETURNING *;
142+
$$ LANGUAGE SQL;
143+
144+
statement ok
145+
CREATE FUNCTION f_fk_p_del(r INT) RETURNS RECORD AS $$
146+
DELETE FROM parent WHERE p = r RETURNING *;
147+
$$ LANGUAGE SQL;
148+
149+
statement ok
150+
CREATE FUNCTION f_fk_c_p_del(k INT, r INT) RETURNS RECORD AS $$
151+
DELETE FROM child WHERE c = k RETURNING *;
152+
DELETE FROM parent WHERE p = r RETURNING *;
153+
$$ LANGUAGE SQL;
154+
155+
statement ok
156+
CREATE FUNCTION f_fk_p_c_del(k INT, r INT) RETURNS RECORD AS $$
157+
DELETE FROM parent WHERE p = r RETURNING *;
158+
DELETE FROM child WHERE c = k RETURNING *;
159+
$$ LANGUAGE SQL;
160+
161+
statement ok
162+
SELECT f_fk_p_del(4);
163+
164+
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"\.
165+
SELECT f_fk_p_del(3);
166+
167+
statement ok
168+
SELECT f_fk_c_del(102);
169+
170+
statement ok
171+
SELECT f_fk_p_del(3);
172+
173+
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"\.
174+
SELECT f_fk_p_c_del(101,2);
175+
176+
statement ok
177+
SELECT f_fk_c_p_del(101,2);
178+
179+
statement ok
180+
SELECT f_fk_c_del(100), f_fk_p_del(1);
181+
182+
query I rowsort
183+
SELECT * FROM parent
184+
----
185+
186+
query II rowsort
187+
SELECT * FROM child
188+
----
189+
190+
191+
subtest upsert
192+
193+
statement ok
194+
TRUNCATE parent CASCADE
195+
196+
statement ok
197+
CREATE FUNCTION f_fk_c_ocdu(k INT, r INT) RETURNS RECORD AS $$
198+
INSERT INTO child VALUES (k, r) ON CONFLICT (c) DO UPDATE SET p = r RETURNING *;
199+
$$ LANGUAGE SQL;
200+
201+
statement ok
202+
INSERT INTO parent VALUES (1), (3);
203+
204+
# Insert
205+
statement ok
206+
SELECT f_fk_c_ocdu(100,1);
207+
208+
# Update to value not in parent fails.
209+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
210+
SELECT f_fk_c_ocdu(100,2);
211+
212+
# Inserting value not in parent fails.
213+
statement error pq: insert on table "child" violates foreign key constraint "child_p_fkey"
214+
SELECT f_fk_c_ocdu(101,2);
215+
216+
statement ok
217+
CREATE FUNCTION f_fk_c_ups(k INT, r INT) RETURNS RECORD AS $$
218+
UPSERT INTO child VALUES (k, r) RETURNING *;
219+
$$ LANGUAGE SQL;
220+
221+
statement ok
222+
SELECT f_fk_c_ups(102,3);
223+
224+
statement error pq: upsert on table "child" violates foreign key constraint "child_p_fkey"
225+
SELECT f_fk_c_ups(102,4);
226+
227+
statement error pq: upsert on table "child" violates foreign key constraint "child_p_fkey"
228+
SELECT f_fk_c_ups(103,4);

‎pkg/sql/logictest/tests/fakedist-disk/generated_test.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/sql/logictest/tests/fakedist-vec-off/generated_test.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/sql/logictest/tests/fakedist/generated_test.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/sql/logictest/tests/local-legacy-schema-changer/generated_test.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/sql/logictest/tests/local-mixed-22.2-23.1/generated_test.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/sql/logictest/tests/local-vec-off/generated_test.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)