Merge pull request #2190 from dolthub/fulghum/set-returning-udfs

Initial support for UDFs with `RETURNS SETOF`

Author: fulghum

server/functions/framework/functions.go

@@ -39,8 +39,8 @@ type FunctionInterface interface {
 	// IsStrict returns whether the function is STRICT, which means if any parameter is NULL, then it returns NULL.
 	// Otherwise, if it's not, the NULL input must be handled by user.
 	IsStrict() bool
-	// IsSRF returns whether the function is set returning function, meaning whether the function returns one or more
-	// rows as a result.
+	// IsSRF returns whether the function is a set returning function, meaning whether the
+	// function returns one or more rows as a result.
 	IsSRF() bool
 	// InternalID returns the ID associated with this function.
 	InternalID() id.Id

server/functions/framework/interpreted_function.go

@@ -38,7 +38,6 @@ type InterpretedFunction struct {
 	Variadic           bool
 	IsNonDeterministic bool
 	Strict             bool
-	SRF                bool
 	Statements         []plpgsql.InterpreterOperation
 }
 
@@ -87,7 +86,12 @@ func (iFunc InterpretedFunction) IsStrict() bool {
 
 // IsSRF implements the interface FunctionInterface.
 func (iFunc InterpretedFunction) IsSRF() bool {
-	return iFunc.SRF
+	switch iFunc.ReturnType.TypCategory {
+	case pgtypes.TypeCategory_CompositeTypes:
+		return true
+	default:
+		return false
+	}
 }
 
 // NonDeterministic implements the interface FunctionInterface.
@@ -107,6 +111,7 @@ func (iFunc InterpretedFunction) QuerySingleReturn(ctx *sql.Context, stack plpgs
 	if err != nil {
 		return nil, err
 	}
+
 	return sql.RunInterpreted(ctx, func(subCtx *sql.Context) (any, error) {
 		sch, rowIter, _, err := stack.Runner().QueryWithBindings(subCtx, stmt, nil, nil, nil)
 		if err != nil {
@@ -164,13 +169,15 @@ func (iFunc InterpretedFunction) QuerySingleReturn(ctx *sql.Context, stack plpgs
 }
 
 // QueryMultiReturn handles queries that may return multiple values over multiple rows.
-func (iFunc InterpretedFunction) QueryMultiReturn(ctx *sql.Context, stack plpgsql.InterpreterStack, stmt string, bindings []string) (rows []sql.Row, err error) {
+func (iFunc InterpretedFunction) QueryMultiReturn(ctx *sql.Context, stack plpgsql.InterpreterStack, stmt string, bindings []string) (schema sql.Schema, rows []sql.Row, err error) {
 	stmt, _, err = iFunc.ApplyBindings(ctx, stack, stmt, bindings, true)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	return sql.RunInterpreted(ctx, func(subCtx *sql.Context) ([]sql.Row, error) {
-		_, rowIter, _, err := stack.Runner().QueryWithBindings(subCtx, stmt, nil, nil, nil)
+
+	rows, err = sql.RunInterpreted(ctx, func(subCtx *sql.Context) (rows []sql.Row, err error) {
+		var rowIter sql.RowIter
+		schema, rowIter, _, err = stack.Runner().QueryWithBindings(subCtx, stmt, nil, nil, nil)
 		if err != nil {
 			return nil, err
 		}
@@ -179,9 +186,10 @@ func (iFunc InterpretedFunction) QueryMultiReturn(ctx *sql.Context, stack plpgsq
 		//  fine.
 		return sql.RowIterToRows(subCtx, rowIter)
 	})
+	return schema, rows, err
 }
 
-// ApplyBindings applies the given bindings to the statement. If `varFound` is false, then the error will be state that
+// ApplyBindings applies the given bindings to the statement. If `varFound` is false, then the error will state that
 // the variable was not found (which means the error may be ignored if you're only concerned with finding a variable).
 // If `varFound` is true, then the error is related to formatting the variable. `enforceType` adds casting and quotes to
 // ensure that the value is correctly represented in the string.

server/plpgsql/interpreter_logic.go

@@ -38,8 +38,11 @@ type InterpretedFunction interface {
 	GetParameterNames() []string
 	GetReturn() *pgtypes.DoltgresType
 	GetStatements() []InterpreterOperation
-	QueryMultiReturn(ctx *sql.Context, stack InterpreterStack, stmt string, bindings []string) (rows []sql.Row, err error)
+	QueryMultiReturn(ctx *sql.Context, stack InterpreterStack, stmt string, bindings []string) (schema sql.Schema, rows []sql.Row, err error)
 	QuerySingleReturn(ctx *sql.Context, stack InterpreterStack, stmt string, targetType *pgtypes.DoltgresType, bindings []string) (val any, err error)
+	// IsSRF returns whether the function is a set returning function, meaning whether the
+	// function returns one or more rows as a result.
+	IsSRF() bool
 }
 
 // GetTypesCollectionFromContext is declared within the core package, but is assigned to this variable to work around
@@ -159,7 +162,7 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 					return nil, err
 				}
 			} else {
-				_, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
+				_, _, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
 				if err != nil {
 					return nil, err
 				}
@@ -200,7 +203,7 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 		case OpCode_InsertInto:
 			// TODO: implement
 		case OpCode_Perform:
-			_, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
+			_, _, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
 			if err != nil {
 				return nil, err
 			}
@@ -229,9 +232,22 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 				sess.Notice(noticeResponse)
 			}
 		case OpCode_Return:
+			// If RETURN QUERY results are being buffered, return those
+			if len(stack.ReturnQueryResults()) > 0 {
+				records := stack.ReturnQueryResults()
+
+				rows := make([]sql.Row, len(records))
+				for i, record := range records {
+					rows[i] = sql.Row{record}
+				}
+
+				return sql.RowsToRowIter(rows...), nil
+			}
+
 			if len(operation.PrimaryData) == 0 {
 				return nil, nil
 			}
+
 			// TODO: handle record types properly, we'll special case triggers for now
 			if iFunc.GetReturn().ID == pgtypes.Trigger.ID && len(operation.SecondaryData) == 1 {
 				normalized := strings.ReplaceAll(strings.ToLower(operation.PrimaryData), " ", "")
@@ -243,7 +259,22 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 					}
 				}
 			}
-			return iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, iFunc.GetReturn(), operation.SecondaryData)
+			val, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, iFunc.GetReturn(), operation.SecondaryData)
+
+			// If this is a set returning function, then we need to return a RowIter and wrap
+			// the composite value in a sql.Row.
+			if iFunc.IsSRF() {
+				return sql.RowsToRowIter(sql.Row{val}), nil
+			}
+			return val, err
+
+		case OpCode_ReturnQuery:
+			schema, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
+			if err != nil {
+				return nil, err
+			}
+			stack.BufferReturnQueryResults(convertRowsToRecords(schema, rows))
+
 		case OpCode_ScopeBegin:
 			stack.PushScope()
 		case OpCode_ScopeEnd:
@@ -259,6 +290,30 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (
 	return nil, nil
 }
 
+// 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 {
+	records := make([][]pgtypes.RecordValue, 0, len(rows))
+	for _, row := range rows {
+		record := make([]pgtypes.RecordValue, len(row))
+		for i, field := range row {
+			t := schema[i].Type
+			doltgresType, ok := t.(*pgtypes.DoltgresType)
+			if !ok {
+				panic("expected Doltgres type")
+			}
+
+			record[i] = pgtypes.RecordValue{
+				Value: field,
+				Type:  doltgresType,
+			}
+		}
+		records = append(records, record)
+	}
+
+	return records
+}
+
 // applyNoticeOptions adds the specified |options| to the |noticeResponse|.
 func applyNoticeOptions(ctx *sql.Context, noticeResponse *pgproto3.NoticeResponse, options map[string]string) error {
 	for key, value := range options {

server/plpgsql/interpreter_operation.go

@@ -19,24 +19,29 @@ package plpgsql
 type OpCode uint16
 
 const (
-	OpCode_Alias      OpCode = iota // https://www.postgresql.org/docs/15/plpgsql-declarations.html#PLPGSQL-DECLARATION-ALIAS
-	OpCode_Assign                   // https://www.postgresql.org/docs/15/plpgsql-statements.html#PLPGSQL-STATEMENTS-ASSIGNMENT
-	OpCode_Case                     // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-CONDITIONALS
-	OpCode_Declare                  // https://www.postgresql.org/docs/15/plpgsql-declarations.html
-	OpCode_DeleteInto               // https://www.postgresql.org/docs/15/plpgsql-statements.html
-	OpCode_Exception                // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-ERROR-TRAPPING
-	OpCode_Execute                  // Executing a standard SQL statement (expects no rows returned unless Target is specified)
-	OpCode_Get                      // https://www.postgresql.org/docs/15/plpgsql-statements.html#PLPGSQL-STATEMENTS-DIAGNOSTICS
-	OpCode_Goto                     // All control-flow structures can be represented using Goto
-	OpCode_If                       // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-CONDITIONALS
-	OpCode_InsertInto               // https://www.postgresql.org/docs/15/plpgsql-statements.html
-	OpCode_Perform                  // https://www.postgresql.org/docs/15/plpgsql-statements.html
-	OpCode_Raise                    // https://www.postgresql.org/docs/15/plpgsql-errors-and-messages.html
-	OpCode_Return                   // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-STATEMENTS-RETURNING
-	OpCode_ScopeBegin               // This is used for scope control, specific to Doltgres
-	OpCode_ScopeEnd                 // This is used for scope control, specific to Doltgres
-	OpCode_SelectInto               // https://www.postgresql.org/docs/15/plpgsql-statements.html
-	OpCode_UpdateInto               // https://www.postgresql.org/docs/15/plpgsql-statements.html
+	// New OpCode values MUST be added to the END of this list!
+	// Function OpCodes are persisted to disk, so these values MUST be stable across Doltgres versions.
+	OpCode_Alias       OpCode = 0  // https://www.postgresql.org/docs/15/plpgsql-declarations.html#PLPGSQL-DECLARATION-ALIAS
+	OpCode_Assign      OpCode = 1  // https://www.postgresql.org/docs/15/plpgsql-statements.html#PLPGSQL-STATEMENTS-ASSIGNMENT
+	OpCode_Case        OpCode = 2  // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-CONDITIONALS
+	OpCode_Declare     OpCode = 3  // https://www.postgresql.org/docs/15/plpgsql-declarations.html
+	OpCode_DeleteInto  OpCode = 4  // https://www.postgresql.org/docs/15/plpgsql-statements.html
+	OpCode_Exception   OpCode = 5  // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-ERROR-TRAPPING
+	OpCode_Execute     OpCode = 6  // Executing a standard SQL statement (expects no rows returned unless Target is specified)
+	OpCode_Get         OpCode = 7  // https://www.postgresql.org/docs/15/plpgsql-statements.html#PLPGSQL-STATEMENTS-DIAGNOSTICS
+	OpCode_Goto        OpCode = 8  // All control-flow structures can be represented using Goto
+	OpCode_If          OpCode = 9  // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-CONDITIONALS
+	OpCode_InsertInto  OpCode = 10 // https://www.postgresql.org/docs/15/plpgsql-statements.html
+	OpCode_Perform     OpCode = 11 // https://www.postgresql.org/docs/15/plpgsql-statements.html
+	OpCode_Raise       OpCode = 12 // https://www.postgresql.org/docs/15/plpgsql-errors-and-messages.html
+	OpCode_Return      OpCode = 13 // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-STATEMENTS-RETURNING
+	OpCode_ScopeBegin  OpCode = 14 // This is used for scope control, specific to Doltgres
+	OpCode_ScopeEnd    OpCode = 15 // This is used for scope control, specific to Doltgres
+	OpCode_SelectInto  OpCode = 16 // https://www.postgresql.org/docs/15/plpgsql-statements.html
+	OpCode_UpdateInto  OpCode = 17 // https://www.postgresql.org/docs/15/plpgsql-statements.html
+	OpCode_ReturnQuery OpCode = 18 // https://www.postgresql.org/docs/current/plpgsql-control-structures.html#PLPGSQL-STATEMENTS-RETURNING-RETURN-NEXT
+	// New OpCode values MUST be added to the END of this list!
+	// Function OpCodes are persisted to disk, so these values MUST be stable across Doltgres versions.
 )
 
 // InterpreterOperation is an operation that will be performed by the interpreter.

server/plpgsql/interpreter_stack.go

@@ -55,6 +55,9 @@ type InterpreterStack struct {
 	stack   *utils.Stack[*InterpreterScopeDetails]
 	runner  sql.StatementRunner
 	labelID int
+
+	// returnQueryBuffer buffers results from RETURN QUERY statements
+	returnQueryBuffer [][]pgtypes.RecordValue
 }
 
 // NewInterpreterStack creates a new InterpreterStack.
@@ -235,3 +238,15 @@ func (is *InterpreterStack) SetAnonymousLabel() {
 	is.stack.Peek().label = fmt.Sprintf("\t%d", is.labelID)
 	is.labelID++
 }
+
+// BufferReturnQueryResults buffers |results| from a RETURN QUERY statement so that they can be returned when
+// the function exits. If results from a previous RETURN QUERY call have already been buffered, |results| will
+// be appended.
+func (is *InterpreterStack) BufferReturnQueryResults(results [][]pgtypes.RecordValue) {
+	is.returnQueryBuffer = append(is.returnQueryBuffer, results...)
+}
+
+// ReturnQueryResults returns the buffered results from a RETURN QUERY statement.
+func (is *InterpreterStack) ReturnQueryResults() [][]pgtypes.RecordValue {
+	return is.returnQueryBuffer
+}

server/plpgsql/json.go

@@ -203,6 +203,12 @@ type plpgSQL_stmt_return struct {
 	LineNumber int32 `json:"lineno"`
 }
 
+// plpgSQL_stmt_return_query exists to match the expected JSON format.
+type plpgSQL_stmt_return_query struct {
+	Query      expr  `json:"query"`
+	LineNumber int32 `json:"lineno"`
+}
+
 // plpgSQL_stmt_while exists to match the expected JSON format.
 type plpgSQL_stmt_while struct {
 	Condition  cond        `json:"cond"`
@@ -231,17 +237,18 @@ type sqlstmt struct {
 // statement exists to match the expected JSON format. Unlike other structs, this is used like a union rather than
 // having a singular expected implementation.
 type statement struct {
-	Assignment *plpgSQL_stmt_assign  `json:"PLpgSQL_stmt_assign"`
-	Case       *plpgSQL_stmt_case    `json:"PLpgSQL_stmt_case"`
-	ExecSQL    *plpgSQL_stmt_execsql `json:"PLpgSQL_stmt_execsql"`
-	Exit       *plpgSQL_stmt_exit    `json:"PLpgSQL_stmt_exit"`
-	If         *plpgSQL_stmt_if      `json:"PLpgSQL_stmt_if"`
-	Loop       *plpgSQL_stmt_loop    `json:"PLpgSQL_stmt_loop"`
-	Perform    *plpgSQL_stmt_perform `json:"PLpgSQL_stmt_perform"`
-	Raise      *plpgSQL_stmt_raise   `json:"PLpgSQL_stmt_raise"`
-	Return     *plpgSQL_stmt_return  `json:"PLpgSQL_stmt_return"`
-	When       *plpgSQL_case_when    `json:"PLpgSQL_case_when"`
-	While      *plpgSQL_stmt_while   `json:"PLpgSQL_stmt_while"`
+	Assignment  *plpgSQL_stmt_assign       `json:"PLpgSQL_stmt_assign"`
+	Case        *plpgSQL_stmt_case         `json:"PLpgSQL_stmt_case"`
+	ExecSQL     *plpgSQL_stmt_execsql      `json:"PLpgSQL_stmt_execsql"`
+	Exit        *plpgSQL_stmt_exit         `json:"PLpgSQL_stmt_exit"`
+	If          *plpgSQL_stmt_if           `json:"PLpgSQL_stmt_if"`
+	Loop        *plpgSQL_stmt_loop         `json:"PLpgSQL_stmt_loop"`
+	Perform     *plpgSQL_stmt_perform      `json:"PLpgSQL_stmt_perform"`
+	Raise       *plpgSQL_stmt_raise        `json:"PLpgSQL_stmt_raise"`
+	Return      *plpgSQL_stmt_return       `json:"PLpgSQL_stmt_return"`
+	ReturnQuery *plpgSQL_stmt_return_query `json:"PLpgSQL_stmt_return_query"`
+	When        *plpgSQL_case_when         `json:"PLpgSQL_case_when"`
+	While       *plpgSQL_stmt_while        `json:"PLpgSQL_stmt_while"`
 }
 
 // Convert converts the JSON statement into its output form.
@@ -522,6 +529,13 @@ func (stmt *plpgSQL_stmt_return) Convert() Return {
 	}
 }
 
+// Convert converts the JSON statement into its output form.
+func (stmt *plpgSQL_stmt_return_query) Convert() ReturnQuery {
+	return ReturnQuery{
+		Query: stmt.Query.Expression.Query,
+	}
+}
+
 // Convert converts the JSON statement into its output form.
 func (stmt *plpgSQL_stmt_while) Convert() (block Block, err error) {
 	// Convert the body of the loop first so we can determine the GOTO offsets

server/plpgsql/json_convert.go

@@ -86,6 +86,8 @@ func jsonConvertStatement(stmt statement) (Statement, error) {
 		return stmt.Raise.Convert(), nil
 	case stmt.Return != nil:
 		return stmt.Return.Convert(), nil
+	case stmt.ReturnQuery != nil:
+		return stmt.ReturnQuery.Convert(), nil
 	case stmt.While != nil:
 		return stmt.While.Convert()
 	default:

server/plpgsql/statements.go

@@ -298,6 +298,33 @@ type Record struct {
 	Fields []string
 }
 
+// ReturnQuery represents a RETURN QUERY statement.
+type ReturnQuery struct {
+	Query string
+}
+
+var _ Statement = ReturnQuery{}
+
+// OperationSize implements the interface Statement.
+func (r ReturnQuery) OperationSize() int32 {
+	return 1
+}
+
+// AppendOperations implements the interface Statement.
+func (r ReturnQuery) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
+	query, referencedVariables, err := substituteVariableReferences(r.Query, stack)
+	if err != nil {
+		return err
+	}
+
+	*ops = append(*ops, InterpreterOperation{
+		OpCode:        OpCode_ReturnQuery,
+		PrimaryData:   query,
+		SecondaryData: referencedVariables,
+	})
+	return nil
+}
+
 // Return represents a RETURN statement.
 type Return struct {
 	Expression string