Skip to content

Commit b9cdf07

Browse files
DrewKimballDrewKimball
committed
plpgsql: add support for ELSIF branches
This patch adds support for executing PLpgSQL `IF` statements with `ELSIF` branches (else if). `IF` statements were already executed as CASE statements under the hood, so this change only requires building the `ELSIF` branches and appending them to the `whens` list. Informs cockroachdb#105254 Release note (sql change): Added support for specifying PLpgSQL `IF` statements with `ELSIF` branches.
1 parent 3ef6592 commit b9cdf07

File tree

5 files changed

+305
-20
lines changed

5 files changed

+305
-20
lines changed

‎pkg/sql/logictest/testdata/logic_test/udf_plpgsql‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,21 @@ CREATE OR REPLACE FUNCTION f(a INT, b INT) RETURNS INT AS $$
365365
END
366366
$$ LANGUAGE PLpgSQL;
367367

368+
statement error pgcode 2F005 control reached end of function without RETURN
369+
CREATE OR REPLACE FUNCTION f(a INT, b INT) RETURNS INT AS $$
370+
DECLARE
371+
i INT;
372+
BEGIN
373+
IF a < b THEN
374+
RETURN -1;
375+
ELSIF a = b THEN
376+
i := 0;
377+
ELSE
378+
RETURN 1;
379+
END IF;
380+
END
381+
$$ LANGUAGE PLpgSQL;
382+
368383
statement error pgcode 2F005 control reached end of function without RETURN
369384
CREATE OR REPLACE FUNCTION f(a INT, b INT) RETURNS INT AS $$
370385
BEGIN
@@ -709,6 +724,7 @@ CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
709724
END
710725
$$ LANGUAGE PLpgSQL;
711726

727+
# Testing CONSTANT variable declarations.
712728
statement ok
713729
CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
714730
DECLARE
@@ -808,3 +824,147 @@ CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
808824
RETURN i;
809825
END
810826
$$ LANGUAGE PLpgSQL;
827+
828+
# Testing IF statements with ELSIF branches.
829+
statement ok
830+
CREATE OR REPLACE FUNCTION f(n INT) RETURNS INT AS $$
831+
DECLARE
832+
i INT := -1;
833+
BEGIN
834+
IF n = 0 THEN
835+
RETURN 0;
836+
ELSIF n = 1 THEN
837+
i := 100;
838+
ELSIF n = 2 THEN
839+
i := 200;
840+
RETURN i;
841+
END IF;
842+
RETURN i;
843+
END
844+
$$ LANGUAGE PLpgSQL;
845+
846+
query IIII
847+
SELECT f(0), f(1), f(2), f(100);
848+
----
849+
0 100 200 -1
850+
851+
statement ok
852+
CREATE OR REPLACE FUNCTION f(a INT, b INT) RETURNS INT AS $$
853+
DECLARE
854+
i INT := 0;
855+
foo INT := 0;
856+
BEGIN
857+
LOOP IF i >= a THEN EXIT; END IF;
858+
IF b = 0 THEN
859+
RETURN NULL;
860+
ELSIF b = 1 THEN
861+
foo := foo + 1;
862+
ELSIF b = 2 THEN
863+
RETURN 100;
864+
ELSE
865+
foo := foo + b;
866+
END IF;
867+
i := i + 1;
868+
END LOOP;
869+
RETURN foo;
870+
END
871+
$$ LANGUAGE PLpgSQL;
872+
873+
query IIIII
874+
SELECT f(0, 0), f(1, 0), f(1, 1), f(1, 2), f(1, 3);
875+
----
876+
0 NULL 1 100 3
877+
878+
query IIII
879+
SELECT f(5, 0), f(5, 1), f(5, 2), f(5, 3);
880+
----
881+
NULL 5 100 15
882+
883+
# Branches should only be executed if the previous ones fail.
884+
statement ok
885+
CREATE OR REPLACE FUNCTION f(n INT) RETURNS INT AS $$
886+
BEGIN
887+
IF n <= 0 THEN
888+
RAISE NOTICE 'foo';
889+
ELSIF n <= 1 THEN
890+
RAISE NOTICE 'bar';
891+
ELSIF n <= 2 THEN
892+
RAISE NOTICE 'baz';
893+
END IF;
894+
RETURN 0;
895+
END
896+
$$ LANGUAGE PLpgSQL;
897+
898+
query T noticetrace
899+
SELECT f(0);
900+
----
901+
NOTICE: foo
902+
903+
query T noticetrace
904+
SELECT f(1);
905+
----
906+
NOTICE: bar
907+
908+
query T noticetrace
909+
SELECT f(2);
910+
----
911+
NOTICE: baz
912+
913+
query I
914+
SELECT f(100);
915+
----
916+
0
917+
918+
# Test nested IF/ELSIF/ELSE.
919+
statement ok
920+
CREATE OR REPLACE FUNCTION f(a INT, b INT) RETURNS INT AS $$
921+
BEGIN
922+
IF a > 1 THEN
923+
IF b > 1 THEN
924+
RETURN 0;
925+
ELSIF b = 1 THEN
926+
RETURN 1;
927+
ELSIF b = 0 THEN
928+
RETURN 2;
929+
ELSE
930+
RETURN 3;
931+
END IF;
932+
ELSIF a = 1 THEN
933+
IF b > 1 THEN
934+
RETURN 4;
935+
ELSIF b = 1 THEN
936+
RETURN 5;
937+
ELSIF b = 0 THEN
938+
RETURN 6;
939+
ELSE
940+
RETURN 7;
941+
END IF;
942+
ELSIF a = 0 THEN
943+
IF b > 1 THEN
944+
RETURN 8;
945+
ELSIF b = 1 THEN
946+
RETURN 9;
947+
ELSIF b = 0 THEN
948+
RETURN 10;
949+
ELSE
950+
RETURN 11;
951+
END IF;
952+
ELSE
953+
IF b > 1 THEN
954+
RETURN 12;
955+
ELSIF b = 1 THEN
956+
RETURN 13;
957+
ELSIF b = 0 THEN
958+
RETURN 14;
959+
ELSE
960+
RETURN 15;
961+
END IF;
962+
END IF;
963+
END
964+
$$ LANGUAGE PLpgSQL;
965+
966+
query IIIIIIIIIIIIIIII
967+
SELECT f(-1, -1), f(-1, 0), f(-1, 1), f(-1, 10), f(0, -1), f(0, 0), f(0, 1), f(0, 10),
968+
f(1, -1), f(1, 0), f(1, 1), f(1, 10), f(10, -1), f(10, 0), f(10, 1), f(10, 10);
969+
----
970+
15 14 13 12 11 10 9 8 7 6 5 4 3 2 1 0

‎pkg/sql/opt/optbuilder/plpgsql.go‎

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -243,12 +243,6 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(
243243
// name as the variable being assigned.
244244
s = b.addPLpgSQLAssign(s, t.Var, t.Value)
245245
case *plpgsqltree.PLpgSQLStmtIf:
246-
if len(t.ElseIfList) != 0 {
247-
panic(unimplemented.New(
248-
"ELSIF statements",
249-
"PL/pgSQL ELSIF branches are not yet supported",
250-
))
251-
}
252246
// IF statement control flow is handled by calling a "continuation"
253247
// function in each branch that executes all the statements that logically
254248
// follow the IF statement block.
@@ -266,25 +260,40 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(
266260
// function at the end of construction in order to resume execution after
267261
// the IF block.
268262
thenScope := b.buildPLpgSQLStatements(t.ThenBody, s.push())
263+
elsifScopes := make([]*scope, len(t.ElseIfList))
264+
for j := range t.ElseIfList {
265+
elsifScopes[j] = b.buildPLpgSQLStatements(t.ElseIfList[j].Stmts, s.push())
266+
}
269267
// Note that if the ELSE body is empty, elseExpr will be equivalent to
270268
// executing the statements following the IF statement (it will be a call
271269
// to the continuation that was built above).
272270
elseScope := b.buildPLpgSQLStatements(t.ElseBody, s.push())
273271
b.popContinuation()
274272

273+
// If one of the branches does not terminate, return nil to indicate a
274+
// non-terminal branch.
275275
if thenScope == nil || elseScope == nil {
276-
// One or both branches didn't terminate with a RETURN statement.
277276
return nil
278277
}
279-
thenExpr, elseExpr := thenScope.expr, elseScope.expr
278+
for j := range elsifScopes {
279+
if elsifScopes[j] == nil {
280+
return nil
281+
}
282+
}
280283

281-
// Build a scalar CASE statement that conditionally executes either branch
284+
// Build a scalar CASE statement that conditionally executes each branch
282285
// of the IF statement as a subquery.
283286
cond := b.buildPLpgSQLExpr(t.Condition, types.Bool, s)
284-
thenScalar := b.ob.factory.ConstructSubquery(thenExpr, &memo.SubqueryPrivate{})
285-
elseScalar := b.ob.factory.ConstructSubquery(elseExpr, &memo.SubqueryPrivate{})
286-
whenExpr := memo.ScalarListExpr{b.ob.factory.ConstructWhen(cond, thenScalar)}
287-
scalar := b.ob.factory.ConstructCase(memo.TrueSingleton, whenExpr, elseScalar)
287+
thenScalar := b.ob.factory.ConstructSubquery(thenScope.expr, &memo.SubqueryPrivate{})
288+
whens := make(memo.ScalarListExpr, 0, len(t.ElseIfList)+1)
289+
whens = append(whens, b.ob.factory.ConstructWhen(cond, thenScalar))
290+
for j := range t.ElseIfList {
291+
elsifCond := b.buildPLpgSQLExpr(t.ElseIfList[j].Condition, types.Bool, s)
292+
elsifScalar := b.ob.factory.ConstructSubquery(elsifScopes[j].expr, &memo.SubqueryPrivate{})
293+
whens = append(whens, b.ob.factory.ConstructWhen(elsifCond, elsifScalar))
294+
}
295+
elseScalar := b.ob.factory.ConstructSubquery(elseScope.expr, &memo.SubqueryPrivate{})
296+
scalar := b.ob.factory.ConstructCase(memo.TrueSingleton, whens, elseScalar)
288297

289298
// Return a single column that projects the result of the CASE statement.
290299
returnColName := scopeColName("").WithMetadataName(b.makeIdentifier("stmt_if"))

‎pkg/sql/opt/optbuilder/testdata/udf_plpgsql‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3079,3 +3079,120 @@ project
30793079
│ └── projections
30803080
│ └── const: 0 [as=stmt_return_3:2]
30813081
└── const: 1
3082+
3083+
# Testing IF statements with ELSIF branches.
3084+
exec-ddl
3085+
CREATE OR REPLACE FUNCTION f(n INT) RETURNS INT AS $$
3086+
DECLARE
3087+
i INT := -1;
3088+
BEGIN
3089+
IF n = 0 THEN
3090+
RETURN 0;
3091+
ELSIF n = 1 THEN
3092+
i := 100;
3093+
ELSIF n = 2 THEN
3094+
i := 200;
3095+
RETURN i;
3096+
END IF;
3097+
RETURN i;
3098+
END
3099+
$$ LANGUAGE PLpgSQL;
3100+
----
3101+
3102+
build format=show-scalars
3103+
SELECT f(0);
3104+
----
3105+
project
3106+
├── columns: f:13
3107+
├── values
3108+
│ └── tuple
3109+
└── projections
3110+
└── udf: f [as=f:13]
3111+
├── args
3112+
│ └── const: 0
3113+
├── params: n:1
3114+
└── body
3115+
└── limit
3116+
├── columns: stmt_if_5:12
3117+
├── project
3118+
│ ├── columns: stmt_if_5:12
3119+
│ ├── project
3120+
│ │ ├── columns: i:2!null
3121+
│ │ ├── values
3122+
│ │ │ └── tuple
3123+
│ │ └── projections
3124+
│ │ └── const: -1 [as=i:2]
3125+
│ └── projections
3126+
│ └── case [as=stmt_if_5:12]
3127+
│ ├── true
3128+
│ ├── when
3129+
│ │ ├── eq
3130+
│ │ │ ├── variable: n:1
3131+
│ │ │ └── const: 0
3132+
│ │ └── subquery
3133+
│ │ └── project
3134+
│ │ ├── columns: stmt_return_3:6!null
3135+
│ │ ├── values
3136+
│ │ │ └── tuple
3137+
│ │ └── projections
3138+
│ │ └── const: 0 [as=stmt_return_3:6]
3139+
│ ├── when
3140+
│ │ ├── eq
3141+
│ │ │ ├── variable: n:1
3142+
│ │ │ └── const: 1
3143+
│ │ └── subquery
3144+
│ │ └── project
3145+
│ │ ├── columns: stmt_if_1:8
3146+
│ │ ├── project
3147+
│ │ │ ├── columns: i:7!null
3148+
│ │ │ ├── values
3149+
│ │ │ │ └── tuple
3150+
│ │ │ └── projections
3151+
│ │ │ └── const: 100 [as=i:7]
3152+
│ │ └── projections
3153+
│ │ └── udf: stmt_if_1 [as=stmt_if_1:8]
3154+
│ │ ├── args
3155+
│ │ │ ├── variable: i:7
3156+
│ │ │ └── variable: n:1
3157+
│ │ ├── params: i:3 n:4
3158+
│ │ └── body
3159+
│ │ └── project
3160+
│ │ ├── columns: stmt_return_2:5
3161+
│ │ ├── values
3162+
│ │ │ └── tuple
3163+
│ │ └── projections
3164+
│ │ └── variable: i:3 [as=stmt_return_2:5]
3165+
│ ├── when
3166+
│ │ ├── eq
3167+
│ │ │ ├── variable: n:1
3168+
│ │ │ └── const: 2
3169+
│ │ └── subquery
3170+
│ │ └── project
3171+
│ │ ├── columns: stmt_return_4:10!null
3172+
│ │ ├── project
3173+
│ │ │ ├── columns: i:9!null
3174+
│ │ │ ├── values
3175+
│ │ │ │ └── tuple
3176+
│ │ │ └── projections
3177+
│ │ │ └── const: 200 [as=i:9]
3178+
│ │ └── projections
3179+
│ │ └── variable: i:9 [as=stmt_return_4:10]
3180+
│ └── subquery
3181+
│ └── project
3182+
│ ├── columns: stmt_if_1:11
3183+
│ ├── values
3184+
│ │ └── tuple
3185+
│ └── projections
3186+
│ └── udf: stmt_if_1 [as=stmt_if_1:11]
3187+
│ ├── args
3188+
│ │ ├── variable: i:2
3189+
│ │ └── variable: n:1
3190+
│ ├── params: i:3 n:4
3191+
│ └── body
3192+
│ └── project
3193+
│ ├── columns: stmt_return_2:5
3194+
│ ├── values
3195+
│ │ └── tuple
3196+
│ └── projections
3197+
│ └── variable: i:3 [as=stmt_return_2:5]
3198+
└── const: 1

‎pkg/sql/plpgsql/parser/plpgsql.y‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,8 @@ func (u *plpgsqlSymUnion) pLpgSQLStmtGetDiagItemList() plpgsqltree.PLpgSQLStmtGe
114114
return u.val.(plpgsqltree.PLpgSQLStmtGetDiagItemList)
115115
}
116116

117-
func (u *plpgsqlSymUnion) pLpgSQLStmtIfElseIfArmList() []*plpgsqltree.PLpgSQLStmtIfElseIfArm {
118-
return u.val.([]*plpgsqltree.PLpgSQLStmtIfElseIfArm)
117+
func (u *plpgsqlSymUnion) pLpgSQLStmtIfElseIfArmList() []plpgsqltree.PLpgSQLStmtIfElseIfArm {
118+
return u.val.([]plpgsqltree.PLpgSQLStmtIfElseIfArm)
119119
}
120120

121121
func (u *plpgsqlSymUnion) pLpgSQLStmtOpen() *plpgsqltree.PLpgSQLStmtOpen {
@@ -315,7 +315,7 @@ func (u *plpgsqlSymUnion) plpgsqlOptionExprs() []plpgsqltree.PLpgSQLStmtRaiseOpt
315315
%type <str> opt_error_level option_type
316316

317317
%type <[]plpgsqltree.PLpgSQLStatement> proc_sect
318-
%type <[]*plpgsqltree.PLpgSQLStmtIfElseIfArm> stmt_elsifs
318+
%type <[]plpgsqltree.PLpgSQLStmtIfElseIfArm> stmt_elsifs
319319
%type <[]plpgsqltree.PLpgSQLStatement> stmt_else loop_body // TODO is this a list of statement?
320320
%type <plpgsqltree.PLpgSQLStatement> pl_block
321321
%type <plpgsqltree.PLpgSQLStatement> proc_stmt
@@ -822,15 +822,15 @@ stmt_if: IF expr_until_then THEN proc_sect stmt_elsifs stmt_else END_IF IF ';'
822822

823823
stmt_elsifs:
824824
{
825-
$$.val = []*plpgsqltree.PLpgSQLStmtIfElseIfArm{};
825+
$$.val = []plpgsqltree.PLpgSQLStmtIfElseIfArm{};
826826
}
827827
| stmt_elsifs ELSIF expr_until_then THEN proc_sect
828828
{
829829
cond, err := plpgsqllex.(*lexer).ParseExpr($3)
830830
if err != nil {
831831
return setErr(plpgsqllex, err)
832832
}
833-
newStmt := &plpgsqltree.PLpgSQLStmtIfElseIfArm{
833+
newStmt := plpgsqltree.PLpgSQLStmtIfElseIfArm{
834834
Condition: cond,
835835
Stmts: $5.plpgsqlStatements(),
836836
}

0 commit comments

Comments
 (0)