diff --git a/server/plpgsql/json.go b/server/plpgsql/json.go index ab98b7ccd5..4e978919a4 100644 --- a/server/plpgsql/json.go +++ b/server/plpgsql/json.go @@ -577,34 +577,42 @@ func (stmt *plpgSQL_stmt_fori) Convert() (block Block, err error) { // Build the loop body: // [0] InitAssign: varName := lower - // [1] If(condition, GotoOffset:2) → jumps to [3] (first body stmt) when true - // [2] ExitGoto → offset=3+bodySize → jumps to ScopeEnd - // [3..3+N-1] body statements (N = bodySize) - // [3+N] IncrAssign: varName := varName +/- step - // [3+N+1] BackGoto → offset=-(3+bodySize) → jumps back to If at [1] + // [1] SkipIncrGoto → offset=2 → jumps to [3], so the first iteration uses lower unchanged + // [2] IncrAssign: varName := varName +/- step (the CONTINUE target) + // [3] If(condition, GotoOffset:2) → jumps to [5] (first body stmt) when true + // [4] ExitGoto → offset=2+bodySize → jumps to ScopeEnd + // [5..5+N-1] body statements (N = bodySize) + // [5+N] BackGoto → offset=-(3+bodySize) → jumps back to IncrAssign at [2] + // + // The increment sits ahead of the condition, rather than at the end of the body, so that CONTINUE has a + // single operation to jump to that both advances the loop and re-tests the condition. // // Because no variables are declared in this block (the loop variable is already // declared by the caller's DECLARE section), ScopeBegin is at M and the // InitAssign is at M+1, so all offsets are consistent. + block.ContinueTargetOffset = 2 block.Body = []Statement{ Assignment{ VariableName: varName, Expression: lowerExpr, }, + Goto{ + Offset: 2, + }, + Assignment{ + VariableName: varName, + Expression: incrExpr, + }, If{ Condition: condition, GotoOffset: 2, }, Goto{ - Offset: 3 + bodySize, + Offset: 2 + bodySize, }, } block.Body = append(block.Body, convertedBody...) block.Body = append(block.Body, - Assignment{ - VariableName: varName, - Expression: incrExpr, - }, Goto{ Offset: -(3 + bodySize), }, @@ -649,6 +657,10 @@ func (stmt *plpgSQL_stmt_fors) Convert() (block Block, err error) { // [1] ForQueryNext – fetch next row into varName, or jump forward by (bodySize+2) to ScopeEnd // [2..2+bodySize-1] body statements // [2+bodySize] Goto back to ForQueryNext: offset = -(1 + bodySize) + // + // A CONTINUE must fetch the next row rather than re-run the query, so it targets the ForQueryNext at [1] + // rather than the ForQueryInit at [0]. + block.ContinueTargetOffset = 1 block.Body = []Statement{ ForQueryInit{CursorName: cursorName, Query: query}, ForQueryNext{CursorName: cursorName, RecordVar: varName, GotoOffset: bodySize + 2}, diff --git a/server/plpgsql/reconcile_labels.go b/server/plpgsql/reconcile_labels.go index 2168a62840..12f063d92a 100644 --- a/server/plpgsql/reconcile_labels.go +++ b/server/plpgsql/reconcile_labels.go @@ -15,11 +15,18 @@ package plpgsql import ( + "strconv" + "github.com/cockroachdb/errors" "github.com/dolthub/doltgresql/utils" ) +// continueTargetOption is the key, within a loop's ScopeBegin Options, holding the offset from that ScopeBegin +// to the loop's next-iteration operation, which is where a CONTINUE for that loop jumps. Block.AppendOperations +// sets it and reconcileLabels removes it, so it never reaches the persisted operations. +const continueTargetOption = "continue_target" + // labelStackItem is the stack item used while reconciling labels. type labelStackItem struct { label string @@ -61,15 +68,29 @@ func reconcileLabels(ops []InterpreterOperation) error { } } case OpCode_ScopeBegin: + // A CONTINUE defaults to the operation after this one, else we'll continually increase the scope. + // Loops that begin with setup operations instead of their next-iteration step say where it is. + start := opIndex + 1 + if offset, ok := operation.Options[continueTargetOption]; ok { + parsedOffset, err := strconv.Atoi(offset) + if err != nil { + return errors.Wrapf(err, "invalid CONTINUE target for scope at operation %d", opIndex) + } + start = opIndex + parsedOffset + } // We'll push the label and loop status to the stack labels.Push(labelStackItem{ label: operation.PrimaryData, - start: opIndex + 1, // We want to go to the operation after this one, else we'll continually increase the scope + start: start, isLoop: len(operation.Target) > 0, }) - // We clear the label and loop status since we only set them for reconciliation + // We clear the label, loop status, and CONTINUE target since we only set them for reconciliation ops[opIndex].PrimaryData = "" ops[opIndex].Target = "" + delete(ops[opIndex].Options, continueTargetOption) + if len(ops[opIndex].Options) == 0 { + ops[opIndex].Options = nil + } case OpCode_ScopeEnd: stackItem := labels.Pop() for gotoIdx, gotoOp := range gotos { diff --git a/server/plpgsql/statements.go b/server/plpgsql/statements.go index ef5d663c4e..39ef1756a5 100644 --- a/server/plpgsql/statements.go +++ b/server/plpgsql/statements.go @@ -16,6 +16,7 @@ package plpgsql import ( "fmt" + "strconv" "github.com/dolthub/go-mysql-server/sql" @@ -71,6 +72,10 @@ type Block struct { Body []Statement Label string IsLoop bool + // ContinueTargetOffset gives the loop's next-iteration operation, where a CONTINUE for this loop jumps, as + // an offset from the body's first operation. It applies only when IsLoop is true, and the zero value suits + // WHILE and plain LOOP, whose bodies begin with that step rather than with loop setup. + ContinueTargetOffset int32 } var _ Statement = Block{} @@ -107,6 +112,7 @@ func (stmt Block) AppendOperations(ops *[]InterpreterOperation, stack *Interpret stmt.Label = stack.GetCurrentLabel() } } + scopeBeginIndex := len(*ops) *ops = append(*ops, InterpreterOperation{ OpCode: OpCode_ScopeBegin, PrimaryData: stmt.Label, @@ -143,6 +149,14 @@ func (stmt Block) AppendOperations(ops *[]InterpreterOperation, stack *Interpret }) } } + if stmt.IsLoop { + // Declarations are already appended, so the body starts at the next operation. reconcileLabels + // resolves this loop's CONTINUE statements through this offset. + continueTarget := len(*ops) + int(stmt.ContinueTargetOffset) + (*ops)[scopeBeginIndex].Options = map[string]string{ + continueTargetOption: strconv.Itoa(continueTarget - scopeBeginIndex), + } + } for _, innerStmt := range stmt.Body { if err := innerStmt.AppendOperations(ops, stack); err != nil { return err diff --git a/testing/go/plpgsql_continue_test.go b/testing/go/plpgsql_continue_test.go new file mode 100644 index 0000000000..b440c8ccb7 --- /dev/null +++ b/testing/go/plpgsql_continue_test.go @@ -0,0 +1,510 @@ +// Copyright 2026 Dolthub, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package _go + +import ( + "testing" + + "github.com/dolthub/go-mysql-server/sql" +) + +// TestPlpgsqlContinue covers CONTINUE (bare, WHEN, and labelled) in every kind of PL/pgSQL loop. All +// expectations were verified against PostgreSQL 16. Every function carries a guard counter that returns a +// sentinel once the loop exceeds its iteration count, so a loop that restarts instead of advancing fails an +// assertion rather than hanging the test. +func TestPlpgsqlContinue(t *testing.T) { + RunScripts(t, []ScriptTest{ + { + Name: "CONTINUE WHEN in an integer FOR loop", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_when() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; n int := 0; guard int := 0; +BEGIN + FOR i IN 1..5 LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + CONTINUE WHEN i % 2 = 0; + n := n + i; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The even values are skipped, so this is 1 + 3 + 5. + Query: `SELECT f_fori_when();`, + Expected: []sql.Row{{9}}, + }, + }, + }, + { + Name: "bare CONTINUE in an integer FOR loop", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_bare() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; n int := 0; guard int := 0; +BEGIN + FOR i IN 1..5 LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + IF i = 3 THEN + CONTINUE; + END IF; + n := n + i; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // Only i = 3 is skipped, so this is 1 + 2 + 4 + 5. + Query: `SELECT f_fori_bare();`, + Expected: []sql.Row{{12}}, + }, + }, + }, + { + Name: "CONTINUE preserves the integer FOR loop variable", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_var() RETURNS text LANGUAGE plpgsql AS $$ +DECLARE i int; acc text := ''; guard int := 0; +BEGIN + FOR i IN 1..4 LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN 'guard'; END IF; + CONTINUE WHEN i = 2; + acc := acc || i::text; + END LOOP; + RETURN acc; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // i = 2 is skipped and the loop variable keeps counting, so this is 1, 3, 4. + Query: `SELECT f_fori_var();`, + Expected: []sql.Row{{"134"}}, + }, + }, + }, + { + Name: "CONTINUE as the last statement of an integer FOR loop body", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_last() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; n int := 0; guard int := 0; +BEGIN + FOR i IN 1..4 LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + n := n + i; + CONTINUE WHEN i > 0; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The CONTINUE always fires, but nothing follows it, so every value is summed. + Query: `SELECT f_fori_last();`, + Expected: []sql.Row{{10}}, + }, + }, + }, + { + Name: "CONTINUE in an integer FOR loop with BY", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_by() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; n int := 0; guard int := 0; +BEGIN + FOR i IN 1..10 BY 3 LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + CONTINUE WHEN i = 4; + n := n + i; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The loop visits 1, 4, 7, 10 and skips 4, so this is 1 + 7 + 10. + Query: `SELECT f_fori_by();`, + Expected: []sql.Row{{18}}, + }, + }, + }, + { + Name: "CONTINUE in a REVERSE integer FOR loop with BY", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_reverse() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; n int := 0; guard int := 0; +BEGIN + FOR i IN REVERSE 10..1 BY 3 LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + CONTINUE WHEN i = 7; + n := n + i; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The loop counts down 10, 7, 4, 1 and skips 7, so this is 10 + 4 + 1. + Query: `SELECT f_fori_reverse();`, + Expected: []sql.Row{{15}}, + }, + }, + }, + { + Name: "CONTINUE and EXIT in the same integer FOR loop", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_exit() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; n int := 0; guard int := 0; +BEGIN + FOR i IN 1..10 LOOP + guard := guard + 1; + IF guard > 25 THEN RETURN -99; END IF; + CONTINUE WHEN i % 2 = 0; + EXIT WHEN i > 6; + n := n + i; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // Even values are skipped and the loop exits at i = 7, so this is 1 + 3 + 5. + Query: `SELECT f_fori_exit();`, + Expected: []sql.Row{{9}}, + }, + }, + }, + { + Name: "bare CONTINUE in a nested integer FOR loop", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_nested() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; j int; n int := 0; guard int := 0; +BEGIN + FOR i IN 1..3 LOOP + FOR j IN 1..3 LOOP + guard := guard + 1; + IF guard > 30 THEN RETURN -99; END IF; + CONTINUE WHEN j = 2; + n := n + 1; + END LOOP; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The inner CONTINUE only affects the inner loop: 3 outer * 2 counted inner. + Query: `SELECT f_fori_nested();`, + Expected: []sql.Row{{6}}, + }, + }, + }, + { + Name: "labelled CONTINUE of an outer integer FOR loop", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fori_label() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; j int; n int := 0; guard int := 0; +BEGIN + <> + FOR i IN 1..3 LOOP + FOR j IN 1..3 LOOP + guard := guard + 1; + IF guard > 30 THEN RETURN -99; END IF; + CONTINUE outer_loop WHEN j = 2; + n := n + (i * 10 + j); + END LOOP; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // Each outer iteration counts j = 1 and then advances i, so this is 11 + 21 + 31. + Query: `SELECT f_fori_label();`, + Expected: []sql.Row{{63}}, + }, + }, + }, + { + Name: "CONTINUE in a WHILE loop", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_while() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int := 0; n int := 0; guard int := 0; +BEGIN + WHILE i < 5 LOOP + i := i + 1; + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + CONTINUE WHEN i = 3; + n := n + i; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // Only i = 3 is skipped, so this is 1 + 2 + 4 + 5. + Query: `SELECT f_while();`, + Expected: []sql.Row{{12}}, + }, + }, + }, + { + Name: "labelled CONTINUE of an outer WHILE loop", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_while_label() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int := 0; j int; n int := 0; guard int := 0; +BEGIN + <> + WHILE i < 3 LOOP + i := i + 1; + FOR j IN 1..3 LOOP + guard := guard + 1; + IF guard > 30 THEN RETURN -99; END IF; + CONTINUE outer_loop WHEN j = 2; + n := n + (i * 10 + j); + END LOOP; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The labelled CONTINUE re-tests the WHILE condition, so this is 11 + 21 + 31. + Query: `SELECT f_while_label();`, + Expected: []sql.Row{{63}}, + }, + }, + }, + { + Name: "CONTINUE in a plain LOOP", + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_loop() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int := 0; n int := 0; guard int := 0; +BEGIN + LOOP + i := i + 1; + EXIT WHEN i > 5; + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + CONTINUE WHEN i = 3; + n := n + i; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // Only i = 3 is skipped, so this is 1 + 2 + 4 + 5. + Query: `SELECT f_loop();`, + Expected: []sql.Row{{12}}, + }, + }, + }, + { + Name: "CONTINUE WHEN in a FOR..IN..SELECT loop", + SetUpScript: []string{ + `CREATE TABLE c1 (id int, val int);`, + `INSERT INTO c1 VALUES (1, 10), (2, 20), (3, 30);`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fors_when() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE r RECORD; n int := 0; guard int := 0; +BEGIN + FOR r IN SELECT id, val FROM c1 ORDER BY id LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + CONTINUE WHEN r.id = 2; + n := n + r.val; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The second row is skipped, so this is 10 + 30. + Query: `SELECT f_fors_when();`, + Expected: []sql.Row{{40}}, + }, + }, + }, + { + Name: "bare CONTINUE in a FOR..IN..SELECT loop", + SetUpScript: []string{ + `CREATE TABLE c2 (id int, val int);`, + `INSERT INTO c2 VALUES (1, 10), (2, 20), (3, 30);`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fors_bare() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE r RECORD; n int := 0; guard int := 0; +BEGIN + FOR r IN SELECT id, val FROM c2 ORDER BY id LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + IF r.id = 3 THEN + CONTINUE; + END IF; + n := n + r.val; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The third row is skipped, so this is 10 + 20. + Query: `SELECT f_fors_bare();`, + Expected: []sql.Row{{30}}, + }, + }, + }, + { + Name: "CONTINUE and EXIT in the same FOR..IN..SELECT loop", + SetUpScript: []string{ + `CREATE TABLE c3 (id int, val int);`, + `INSERT INTO c3 VALUES (1, 10), (2, 20), (3, 30), (4, 40), (5, 50);`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fors_exit() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE r RECORD; n int := 0; guard int := 0; +BEGIN + FOR r IN SELECT id, val FROM c3 ORDER BY id LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + CONTINUE WHEN r.id = 2; + EXIT WHEN r.id = 4; + n := n + r.val; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The second row is skipped and the loop exits on the fourth, so this is 10 + 30. + Query: `SELECT f_fors_exit();`, + Expected: []sql.Row{{40}}, + }, + }, + }, + { + Name: "labelled CONTINUE of an outer FOR..IN..SELECT loop", + SetUpScript: []string{ + `CREATE TABLE c4 (id int, val int);`, + `INSERT INTO c4 VALUES (1, 10), (2, 20), (3, 30);`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fors_label() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE r RECORD; s RECORD; n int := 0; guard int := 0; +BEGIN + <> + FOR r IN SELECT id, val FROM c4 ORDER BY id LOOP + FOR s IN SELECT id FROM c4 ORDER BY id LOOP + guard := guard + 1; + IF guard > 30 THEN RETURN -99; END IF; + CONTINUE outer_loop WHEN s.id = 2; + n := n + (r.id * 10 + s.id); + END LOOP; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // Each outer row counts its first inner row, so this is 11 + 21 + 31. + Query: `SELECT f_fors_label();`, + Expected: []sql.Row{{63}}, + }, + }, + }, + { + Name: "labelled CONTINUE of an outer integer FOR loop from a FOR..IN..SELECT loop", + SetUpScript: []string{ + `CREATE TABLE c5 (id int);`, + `INSERT INTO c5 VALUES (1), (2), (3);`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_mixed_label() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE i int; r RECORD; n int := 0; guard int := 0; +BEGIN + <> + FOR i IN 1..3 LOOP + FOR r IN SELECT id FROM c5 ORDER BY id LOOP + guard := guard + 1; + IF guard > 30 THEN RETURN -99; END IF; + CONTINUE outer_loop WHEN r.id = 2; + n := n + (i * 10 + r.id); + END LOOP; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The labelled CONTINUE advances the integer FOR loop, so this is 11 + 21 + 31. + Query: `SELECT f_mixed_label();`, + Expected: []sql.Row{{63}}, + }, + }, + }, + { + Name: "CONTINUE as the last statement of a FOR..IN..SELECT loop body", + SetUpScript: []string{ + `CREATE TABLE c6 (id int, val int);`, + `INSERT INTO c6 VALUES (1, 10), (2, 20), (3, 30);`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_fors_last() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE r RECORD; n int := 0; guard int := 0; +BEGIN + FOR r IN SELECT id, val FROM c6 ORDER BY id LOOP + guard := guard + 1; + IF guard > 20 THEN RETURN -99; END IF; + n := n + r.val; + CONTINUE WHEN r.id > 0; + END LOOP; + RETURN n; +END; $$;`, + Expected: []sql.Row{}, + }, + { + // The CONTINUE always fires, but nothing follows it, so every row is summed. + Query: `SELECT f_fors_last();`, + Expected: []sql.Row{{60}}, + }, + }, + }, + }) +}