Support the built-in FOUND variable in plpgsql

server/plpgsql/interpreter_logic.go

@@ -22,6 +22,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
 	"github.com/dolthub/go-mysql-server/sql"
+	gmstypes "github.com/dolthub/go-mysql-server/sql/types"
 	"github.com/jackc/pgx/v5/pgproto3"
 
 	"github.com/dolthub/doltgresql/core/id"
@@ -217,11 +218,21 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 						return nil, err
 					}
 				}
+				if setsFound(operation) {
+					if err = stack.SetFound(ctx, rowFound); err != nil {
+						return nil, err
+					}
+				}
 			} else {
-				_, _, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
+				_, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
 				if err != nil {
 					return nil, err
 				}
+				if setsFound(operation) {
+					if err = stack.SetFound(ctx, queryProducedRow(rows)); err != nil {
+						return nil, err
+					}
+				}
 			}
 		case OpCode_ExecuteInto:
 			// The target is a RECORD, which takes on the shape of the query's result columns.
@@ -238,26 +249,43 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 			if err = stack.UpdateRecord(operation.Target, schema, row); err != nil {
 				return nil, err
 			}
+			if setsFound(operation) {
+				if err = stack.SetFound(ctx, len(rows) > 0); err != nil {
+					return nil, err
+				}
+			}
 		case OpCode_DeclareRecord:
 			stack.NewRecord(operation.Target, nil, nil)
 		case OpCode_Get:
 			// TODO: implement
 		case OpCode_Goto:
 			// We must compare to the index - 1, so that the increment hits our target
 			if counter <= operation.Index {
+				// Jumping forward leaves every scope it passes over for good, so each one is torn down
+				// the same way reaching its ScopeEnd would tear it down. A labelled EXIT out of a nested
+				// loop passes over that loop's ScopeEnd this way.
 				for ; counter < operation.Index-1; counter++ {
 					switch statements[counter].OpCode {
 					case OpCode_ScopeBegin:
 						stack.PushScope()
 					case OpCode_ScopeEnd:
-						stack.PopScope()
+						if err := exitScope(ctx, stack); err != nil {
+							return nil, err
+						}
 					}
 				}
 			} else {
+				// Jumping backward passes a scope's ScopeBegin only when that scope is being left for
+				// good, so that is torn down like any other scope exit. A labelled CONTINUE of an outer
+				// loop leaves the inner loop this way. Reaching a ScopeEnd backwards is the opposite: a
+				// scope that already closed is being re-entered, and the fresh scope this pushes carries
+				// nothing to tear down when the walk reaches its ScopeBegin a moment later.
 				for ; counter > operation.Index-1; counter-- {
 					switch statements[counter].OpCode {
 					case OpCode_ScopeBegin:
-						stack.PopScope()
+						if err := exitScope(ctx, stack); err != nil {
+							return nil, err
+						}
 					case OpCode_ScopeEnd:
 						stack.PushScope()
 					}
@@ -268,18 +296,27 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 			if err != nil {
 				return nil, err
 			}
-			if retVal.(bool) {
+			conditionMet := retVal.(bool)
+			if isLoopCondition(operation) {
+				// An integer FOR loop has no cursor to carry the fact that its body ran, so its condition
+				// is what records it, for the FOUND the loop reports once it is left.
+				stack.MarkScopeLoop(conditionMet)
+			}
+			if conditionMet {
 				// We're never changing the scope, so we can just assign it directly.
 				// Also, we must assign to index-1, so that the increment hits our target.
 				counter = operation.Index - 1
 			}
 		case OpCode_InsertInto:
 			// TODO: implement
 		case OpCode_Perform:
-			_, _, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
+			_, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
 			if err != nil {
 				return nil, err
 			}
+			if err = stack.SetFound(ctx, queryProducedRow(rows)); err != nil {
+				return nil, err
+			}
 		case OpCode_Raise:
 			// TODO: Use the client_min_messages config param to determine which
 			//       notice levels to send to the client.
@@ -360,13 +397,17 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 				return nil, err
 			}
 			stack.InitCursor(operation.Target, schema, rows)
+			// The loop reports FOUND when it is left even if the query matched nothing.
+			stack.MarkScopeLoop(false)
 		case OpCode_ForQueryNext:
 			schema, row, ok := stack.AdvanceCursor(operation.PrimaryData)
 			if !ok {
-				stack.CloseCursor(operation.PrimaryData)
-				// Jump forward past the loop body and back-goto, same mechanism as OpCode_If.
+				// Jump forward past the loop body and back-goto, same mechanism as OpCode_If. The loop's
+				// ScopeEnd is what closes the cursor and reports FOUND, since every way out of the loop
+				// reaches it and this one does not.
 				counter = operation.Index - 1
 			} else {
+				stack.MarkScopeLoop(true)
 				if err := stack.UpdateRecord(operation.Target, schema, row); err != nil {
 					return nil, err
 				}
@@ -381,10 +422,15 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 				return nil, err
 			}
 			stack.BufferReturnQueryResults(records)
+			if err = stack.SetFound(ctx, len(rows) > 0); err != nil {
+				return nil, err
+			}
 		case OpCode_ScopeBegin:
 			stack.PushScope()
 		case OpCode_ScopeEnd:
-			stack.PopScope()
+			if err := exitScope(ctx, stack); err != nil {
+				return nil, err
+			}
 		case OpCode_SelectInto:
 			// TODO: implement
 		case OpCode_UpdateInto:
@@ -396,6 +442,49 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 	return nil, nil
 }
 
+// exitScope performs everything that leaving a scope entails. Both the ScopeEnd opcode and the forward walk
+// of a Goto leave scopes, so they share this rather than each handling scope depth on its own, which would
+// leave whichever of them grew a new responsibility last out of step with the other.
+//
+// A FOR..IN..SELECT loop's scope owns the cursor the loop iterates, so leaving the scope is what closes it.
+// Leaving it is also what reports a FOR loop's FOUND: Postgres sets FOUND when such a loop exits, by
+// whichever path, to whether the body ran at all, and leaves it alone while the loop is running. A WHILE or
+// plain LOOP does not set FOUND, so only a scope its loop marked reports one.
+func exitScope(ctx *sql.Context, stack InterpreterStack) error {
+	if cursorName := stack.ScopeCursor(); len(cursorName) > 0 {
+		stack.CloseCursor(cursorName)
+	}
+	if reportsFound, iterated := stack.ScopeLoop(); reportsFound {
+		if err := stack.SetFound(ctx, iterated); err != nil {
+			return err
+		}
+	}
+	stack.PopScope()
+	return nil
+}
+
+// setsFound reports whether the operation should update the built-in FOUND variable. Only the opcodes that
+// static and dynamic execution share have to ask: PostgreSQL defines a static statement as setting FOUND
+// and a dynamic EXECUTE as leaving it alone.
+func setsFound(operation InterpreterOperation) bool {
+	return operation.Options[OptionSetsFound] == "true"
+}
+
+// isLoopCondition reports whether the operation is the conditional jump that advances an integer FOR loop.
+func isLoopCondition(operation InterpreterOperation) bool {
+	return operation.Options[OptionLoopCondition] == "true"
+}
+
+// queryProducedRow reports whether a statement produced a row for the purposes of FOUND. A data-modifying
+// statement without RETURNING reports a single OkResult row no matter how many rows it touched, so for those
+// it is the affected-row count that answers the question.
+func queryProducedRow(rows []sql.Row) bool {
+	if len(rows) == 1 && gmstypes.IsOkResult(rows[0]) {
+		return rows[0][0].(gmstypes.OkResult).RowsAffected > 0
+	}
+	return len(rows) > 0
+}
+
 // convertRowsToRecords iterates overs |rows| and converts each field in each row
 // into a RecordValue. |schema| is specified for type information.
 func convertRowsToRecords(schema sql.Schema, rows []sql.Row) ([][]pgtypes.RecordValue, error) {

server/plpgsql/interpreter_operation.go

@@ -48,6 +48,17 @@ const (
 	// Function OpCodes are persisted to disk, so these values MUST be stable across Doltgres versions.
 )
 
+// OptionSetsFound is an Options key marking an operation that updates the built-in FOUND variable. Static
+// and dynamic execution share an opcode, and PostgreSQL defines a static statement as setting FOUND while a
+// dynamic EXECUTE deliberately leaves it alone, so the two are told apart by this option rather than by
+// their opcode. Operations that always set FOUND do not carry it.
+const OptionSetsFound = "sets_found"
+
+// OptionLoopCondition is an Options key marking the conditional jump that advances an integer FOR loop.
+// Every kind of loop compiles to the same conditional jump, and PostgreSQL defines a FOR loop as setting
+// FOUND on exit while a WHILE or a plain LOOP leaves it alone, so the two are told apart by this option.
+const OptionLoopCondition = "loop_condition"
+
 // InterpreterOperation is an operation that will be performed by the interpreter.
 type InterpreterOperation struct {
 	OpCode        OpCode

server/plpgsql/interpreter_stack.go

@@ -100,6 +100,11 @@ const (
 	TriggerOldRecordName = "old"
 )
 
+// FoundVariableName is the name of the built-in FOUND variable, which PL/pgSQL creates for every function
+// and which reports whether the most recent statement that is defined to set it produced any row.
+// https://www.postgresql.org/docs/15/plpgsql-statements.html#PLPGSQL-STATEMENTS-DIAGNOSTICS
+const FoundVariableName = "found"
+
 // cursorState holds the result set for a FOR record IN query LOOP cursor.
 type cursorState struct {
 	Schema sql.Schema
@@ -132,6 +137,14 @@ type InterpreterVariableReference struct {
 type InterpreterScopeDetails struct {
 	variables map[string]*interpreterVariable
 	label     string
+	// cursor names the FOR..IN..SELECT cursor this scope owns, if it is such a loop's scope. The scope
+	// owning it is what lets the cursor be torn down wherever the loop is left, rather than only where
+	// the cursor runs out.
+	cursor string
+	// reportsFound marks the scope of a loop that sets FOUND when it is left, and iterated records
+	// whether that loop ever advanced into its body.
+	reportsFound bool
+	iterated     bool
 }
 
 // InterpreterStack represents the working information that an interpreter will use during execution. It is not exactly
@@ -407,13 +420,21 @@ func (is *InterpreterStack) ReturnOutParamResults() any {
 	return record
 }
 
-// InitCursor stores the result set for a FOR record IN query LOOP cursor.
+// InitCursor stores the result set for a FOR record IN query LOOP cursor. The cursor is opened in the
+// loop's own scope, which takes ownership of it.
 func (is *InterpreterStack) InitCursor(name string, schema sql.Schema, rows []sql.Row) {
 	is.cursors[name] = &cursorState{
 		Schema: schema,
 		Rows:   rows,
 		Index:  0,
 	}
+	is.stack.Peek().cursor = name
+}
+
+// ScopeCursor returns the name of the cursor the current scope owns, or an empty string when the scope is
+// not that of a FOR..IN..SELECT loop.
+func (is *InterpreterStack) ScopeCursor() string {
+	return is.stack.Peek().cursor
 }
 
 // AdvanceCursor returns the next row for the named cursor and advances its index.
@@ -433,6 +454,32 @@ func (is *InterpreterStack) CloseCursor(name string) {
 	delete(is.cursors, name)
 }
 
+// MarkScopeLoop marks the current scope as a loop's, whose exit reports FOUND, and records whether the loop
+// has just advanced into its body. Only the loop itself knows that it advanced, and it cannot record it in
+// FOUND, which PostgreSQL leaves alone until the loop is left.
+func (is *InterpreterStack) MarkScopeLoop(advanced bool) {
+	details := is.stack.Peek()
+	details.reportsFound = true
+	details.iterated = details.iterated || advanced
+}
+
+// ScopeLoop reports whether leaving the current scope sets FOUND, and if so whether its loop body ran.
+func (is *InterpreterStack) ScopeLoop() (reportsFound bool, iterated bool) {
+	details := is.stack.Peek()
+	return details.reportsFound, details.iterated
+}
+
+// SetFound updates the built-in FOUND variable. Functions compiled before FOUND was supported do not declare
+// it, and a function may shadow the name with a variable of its own, so anything other than the built-in
+// boolean is left alone rather than being overwritten.
+func (is *InterpreterStack) SetFound(ctx *sql.Context, found bool) error {
+	ref := is.GetVariable(FoundVariableName)
+	if ref.Type == nil || ref.Type.ID != pgtypes.Bool.ID {
+		return nil
+	}
+	return is.SetVariable(ctx, FoundVariableName, found)
+}
+
 // UpdateRecord finds the named variable and sets its schema and row value. A nil |val| gives the record the
 // shape of |schema| with every field NULL, which is what PostgreSQL does when an INTO query matches no rows.
 func (is *InterpreterStack) UpdateRecord(name string, schema sql.Schema, val sql.Row) error {

server/plpgsql/json.go

@@ -469,6 +469,10 @@ func (stmt *plpgSQL_stmt_execsql) Convert() (ExecuteSQL, error) {
 		Statement:      stmt.SQLStmt.Expr.Query,
 		Target:         target,
 		TargetIsRecord: targetIsRecord,
+		// An INTO clause always reports whether it found a row. Without one, only a data-modifying
+		// statement touches FOUND. The INTO check comes first because the raw text of a SELECT INTO
+		// still carries the INTO clause, which parses as something else entirely.
+		SetsFound: stmt.Into || isDataModifying(stmt.SQLStmt.Expr.Query),
 	}, nil
 }
 
@@ -604,8 +608,9 @@ func (stmt *plpgSQL_stmt_fori) Convert() (block Block, err error) {
 			Expression:   incrExpr,
 		},
 		If{
-			Condition:  condition,
-			GotoOffset: 2,
+			Condition:       condition,
+			GotoOffset:      2,
+			IsLoopCondition: true,
 		},
 		Goto{
 			Offset: 2 + bodySize,

server/plpgsql/json_convert.go

@@ -38,6 +38,11 @@ func jsonConvert(jsonBlock plpgSQL_block) (Block, error) {
 		}
 	}
 	offset := int32(0) - lowestRecordNumber
+	// PL/pgSQL creates the built-in FOUND variable itself, immediately after the function's parameters, so
+	// it arrives looking like a parameter (no line number). We track where it lands so it can be turned
+	// back into a declaration below. A parameter may also be named `found`, in which case the built-in is
+	// the later of the two and shadows it, so the last match is the one that counts.
+	foundVariableIndex := -1
 	// Then we do a second loop that actually adds all of the datums to the block
 	for _, v := range jsonBlock.Datums {
 		switch {
@@ -62,6 +67,9 @@ func jsonConvert(jsonBlock plpgSQL_block) (Block, error) {
 				block.Records[recordParentNumber].Fields, v.RecordField.FieldName)
 		case v.Row != nil:
 		case v.Variable != nil:
+			if v.Variable.LineNumber == 0 && strings.EqualFold(v.Variable.RefName, FoundVariableName) {
+				foundVariableIndex = len(block.Variables)
+			}
 			block.Variables = append(block.Variables, Variable{
 				Name:        v.Variable.RefName,
 				Type:        strings.ToLower(v.Variable.Type.Type.Name),
@@ -74,6 +82,18 @@ func jsonConvert(jsonBlock plpgSQL_block) (Block, error) {
 			return Block{}, errors.New("unhandled declared variable: expected a record, record field, row, or variable")
 		}
 	}
+	// FOUND is not passed in by the caller, so it has to be declared and initialized like any other local.
+	// Postgres starts it out false.
+	if foundVariableIndex >= 0 {
+		block.Variables[foundVariableIndex].IsParameter = false
+		block.Variables[foundVariableIndex].Default = "false"
+		// The datum names the type as `pg_catalog."boolean"`, which OpCode_Declare cannot resolve: it maps
+		// a qualified pg_catalog name through TypeForNonKeywordTypeName, which has no entry for the
+		// `boolean` spelling, so the lookup is left asking for a type named `boolean` and fails. Name the
+		// type by its canonical name instead. The same gap makes a hand-written `DECLARE b
+		// pg_catalog.boolean` fail, so teaching that alias table about `boolean` would let this go away.
+		block.Variables[foundVariableIndex].Type = "pg_catalog.bool"
+	}
 	// The NEW and OLD records of a trigger appear in the datum list like any other record, but they are
 	// supplied by the trigger invocation rather than declared by the function, so they must not be
 	// redeclared (which would shadow the supplied values with an empty record).