diff --git a/server/functions/framework/interpreted_function.go b/server/functions/framework/interpreted_function.go index 0cd55ae455..59b8b7e1fe 100644 --- a/server/functions/framework/interpreted_function.go +++ b/server/functions/framework/interpreted_function.go @@ -195,45 +195,84 @@ func (iFunc InterpretedFunction) QuerySingleReturn(ctx *sql.Context, stack plpgs if len(rows[0]) != 1 { return nil, errors.New("expression returned multiple results") } - if targetType == nil { - return rows[0][0], nil - } - if rows[0][0] == nil { - return nil, nil - } - sourceType, ok := sch[0].Type.(*pgtypes.DoltgresType) - if !ok { - // TODO: We ensure we have a DoltgresType, but we should also convert the value to - // ensure it's in the correct form for the DoltgresType. This logic lives in - // pgexpressions.GMSCast, but need to be extracted to avoid a dependency cycle - // so it can be used here and from server.plpgsql. - sourceType, err = pgtypes.FromGmsTypeToDoltgresType(sch[0].Type) - if err != nil { - return nil, err - } - } - castsColl, err := core.GetCastsCollectionFromContext(ctx, "") + return castQueryValue(subCtx, rows[0][0], sch[0].Type, targetType) + }) +} + +// castQueryValue converts |val|, which a query produced with the column type |columnType|, into the form +// expected by |targetType| using an assignment cast. A nil |targetType| leaves the value as it is. +func castQueryValue(ctx *sql.Context, val any, columnType sql.Type, targetType *pgtypes.DoltgresType) (any, error) { + if targetType == nil { + return val, nil + } + if val == nil { + return nil, nil + } + sourceType, ok := columnType.(*pgtypes.DoltgresType) + if !ok { + // TODO: We ensure we have a DoltgresType, but we should also convert the value to + // ensure it's in the correct form for the DoltgresType. This logic lives in + // pgexpressions.GMSCast, but need to be extracted to avoid a dependency cycle + // so it can be used here and from server.plpgsql. + var err error + sourceType, err = pgtypes.FromGmsTypeToDoltgresType(columnType) if err != nil { return nil, err } - cast, err := castsColl.GetAssignmentCast(ctx, sourceType, targetType) - if err != nil { - return nil, err + } + castsColl, err := core.GetCastsCollectionFromContext(ctx, "") + if err != nil { + return nil, err + } + cast, err := castsColl.GetAssignmentCast(ctx, sourceType, targetType) + if err != nil { + return nil, err + } + if !cast.ID.IsValid() { + // TODO: We're using assignment casting, but for some reason we have to use I/O casting here, which is incorrect? + // We need to dig into this and figure out exactly what's happening, as this is "wrong" according to what + // I understand. This lines up more with explicit casting, but it's supposed to be assignment. + // Maybe there are specific rules for pgsql? + if sourceType.TypCategory == pgtypes.TypeCategory_StringTypes { + cast.ID = id.NewCast(sourceType.ID, targetType.ID) + cast.UseInOut = true + } else { + return nil, errors.New("no valid cast for return value") } - if !cast.ID.IsValid() { - // TODO: We're using assignment casting, but for some reason we have to use I/O casting here, which is incorrect? - // We need to dig into this and figure out exactly what's happening, as this is "wrong" according to what - // I understand. This lines up more with explicit casting, but it's supposed to be assignment. - // Maybe there are specific rules for pgsql? - if sourceType.TypCategory == pgtypes.TypeCategory_StringTypes { - cast.ID = id.NewCast(sourceType.ID, targetType.ID) - cast.UseInOut = true - } else { - return nil, errors.New("no valid cast for return value") + } + return cast.Eval(ctx, val, sourceType, targetType) +} + +// QueryRowReturn handles a statement whose result is written into an INTO clause's targets. It returns the +// first row of the result with each value cast to the matching entry of |targetTypes|, and reports whether +// the query produced a row at all. Any rows after the first are discarded, and no rows is not an error: +// that is what PL/pgSQL does for an INTO clause without STRICT. +func (iFunc InterpretedFunction) QueryRowReturn(ctx *sql.Context, stack plpgsql.InterpreterStack, stmt string, targetTypes []*pgtypes.DoltgresType, bindings []string) (row sql.Row, ok bool, err error) { + schema, rows, err := iFunc.QueryMultiReturn(ctx, stack, stmt, bindings) + if err != nil { + return nil, false, err + } + if len(rows) == 0 { + return nil, false, nil + } + if len(rows[0]) != len(targetTypes) { + return nil, false, errors.Errorf("INTO expects %d values, query returned %d", len(targetTypes), len(rows[0])) + } + row, err = sql.RunInterpreted(ctx, func(subCtx *sql.Context) (sql.Row, error) { + castRow := make(sql.Row, len(targetTypes)) + for i := range castRow { + castVal, err := castQueryValue(subCtx, rows[0][i], schema[i].Type, targetTypes[i]) + if err != nil { + return nil, err } + castRow[i] = castVal } - return cast.Eval(subCtx, rows[0][0], sourceType, targetType) + return castRow, nil }) + if err != nil { + return nil, false, err + } + return row, true, nil } // QueryMultiReturn handles queries that may return multiple values over multiple rows. diff --git a/server/plpgsql/interpreter_logic.go b/server/plpgsql/interpreter_logic.go index eefb1d4314..992890a692 100644 --- a/server/plpgsql/interpreter_logic.go +++ b/server/plpgsql/interpreter_logic.go @@ -40,6 +40,7 @@ type InterpretedFunction interface { GetReturn() *pgtypes.DoltgresType GetStatements() []InterpreterOperation QueryMultiReturn(ctx *sql.Context, stack InterpreterStack, stmt string, bindings []string) (schema sql.Schema, rows []sql.Row, err error) + QueryRowReturn(ctx *sql.Context, stack InterpreterStack, stmt string, targetTypes []*pgtypes.DoltgresType, bindings []string) (row sql.Row, ok bool, 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. @@ -193,43 +194,26 @@ func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) ( // TODO: implement case OpCode_Execute: if len(operation.Target) > 0 { - if vars := strings.Split(operation.Target, ","); len(vars) > 1 { - // multiple column row result - sch, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData) - if err != nil { - return nil, err - } - if len(rows) > 1 { - return nil, errors.New("query returned more than one row") - } - for i, row := range rows { - if len(row) != len(vars) { - return nil, errors.New("number of row values does not match number of schema columns") - } - target := stack.GetVariable(vars[i]) - if target.Type == nil { - return nil, fmt.Errorf("variable `%s` could not be found", operation.Target) - } - if sch[i].Type.(*pgtypes.DoltgresType).ID != target.Type.ID { - return nil, fmt.Errorf("variable type `%s` does not match `%s`", sch[i].Type.String(), target.Type.String()) - } - err = stack.SetVariable(ctx, vars[i], rows[0][i]) - if err != nil { - return nil, err - } - } - } else { - // single column - target := stack.GetVariable(operation.Target) + vars := strings.Split(operation.Target, ",") + targetTypes := make([]*pgtypes.DoltgresType, len(vars)) + for i, varName := range vars { + target := stack.GetVariable(varName) if target.Type == nil { - return nil, fmt.Errorf("variable `%s` could not be found", operation.Target) + return nil, fmt.Errorf("variable `%s` could not be found", varName) } - retVal, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, target.Type, operation.SecondaryData) - if err != nil { - return nil, err + targetTypes[i] = target.Type + } + row, rowFound, err := iFunc.QueryRowReturn(ctx, stack, operation.PrimaryData, targetTypes, operation.SecondaryData) + if err != nil { + return nil, err + } + // When the query matches nothing, every target is set to NULL rather than left alone. + for i, varName := range vars { + var val any + if rowFound { + val = row[i] } - err = stack.SetVariable(ctx, operation.Target, retVal) - if err != nil { + if err = stack.SetVariable(ctx, varName, val); err != nil { return nil, err } } diff --git a/testing/go/plpgsql_into_test.go b/testing/go/plpgsql_into_test.go new file mode 100644 index 0000000000..ffee4bcc9e --- /dev/null +++ b/testing/go/plpgsql_into_test.go @@ -0,0 +1,124 @@ +// 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" +) + +// TestPlpgsqlSelectInto covers what an INTO clause does to its targets when the query returns something +// other than exactly one row, and when it has several targets. All expectations were verified against +// PostgreSQL 16. +func TestPlpgsqlSelectInto(t *testing.T) { + RunScripts(t, []ScriptTest{ + { + Name: "INTO targets when the query matches no rows", + SetUpScript: []string{ + `CREATE TABLE k (id int, nm text);`, + `INSERT INTO k VALUES (1, 'a'), (2, 'b');`, + }, + Assertions: []ScriptTestAssertion{ + { + // Without STRICT this is not an error: the target is set to NULL. + Query: `CREATE FUNCTION f_scalar_miss() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE v int; +BEGIN + SELECT id INTO v FROM k WHERE id = 99; + RETURN coalesce(v, -1); +END; $$;`, + Expected: []sql.Row{}, + }, + { + Query: `SELECT f_scalar_miss();`, + Expected: []sql.Row{{-1}}, + }, + { + // Every target is set to NULL, not just the first. + Query: `CREATE FUNCTION f_two_miss() RETURNS text LANGUAGE plpgsql AS $$ +DECLARE a int; b text; +BEGIN + SELECT id, nm INTO a, b FROM k WHERE id = 99; + RETURN coalesce(a::text, 'NULLa') || '/' || coalesce(b, 'NULLb'); +END; $$;`, + Expected: []sql.Row{}, + }, + { + Query: `SELECT f_two_miss();`, + Expected: []sql.Row{{"NULLa/NULLb"}}, + }, + }, + }, + { + Name: "INTO with several targets assigns every one of them", + SetUpScript: []string{ + `CREATE TABLE k (id int, nm text);`, + `INSERT INTO k VALUES (1, 'a'), (2, 'b');`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `CREATE FUNCTION f_two() RETURNS text LANGUAGE plpgsql AS $$ +DECLARE a int; b text; +BEGIN + SELECT id, nm INTO a, b FROM k WHERE id = 1; + RETURN coalesce(a::text, 'NULLa') || '/' || coalesce(b, 'NULLb'); +END; $$;`, + Expected: []sql.Row{}, + }, + { + Query: `SELECT f_two();`, + Expected: []sql.Row{{"1/a"}}, + }, + }, + }, + { + Name: "INTO keeps the first row when the query matches several", + SetUpScript: []string{ + `CREATE TABLE k (id int, nm text);`, + `INSERT INTO k VALUES (1, 'a'), (2, 'b');`, + }, + Assertions: []ScriptTestAssertion{ + { + // Without STRICT, extra rows are discarded rather than raising. + Query: `CREATE FUNCTION f_scalar_multi() RETURNS int LANGUAGE plpgsql AS $$ +DECLARE v int; +BEGIN + SELECT id INTO v FROM k ORDER BY id; + RETURN v; +END; $$;`, + Expected: []sql.Row{}, + }, + { + Query: `SELECT f_scalar_multi();`, + Expected: []sql.Row{{1}}, + }, + { + Query: `CREATE FUNCTION f_two_multi() RETURNS text LANGUAGE plpgsql AS $$ +DECLARE a int; b text; +BEGIN + SELECT id, nm INTO a, b FROM k ORDER BY id; + RETURN coalesce(a::text, 'NULLa') || '/' || coalesce(b, 'NULLb'); +END; $$;`, + Expected: []sql.Row{}, + }, + { + Query: `SELECT f_two_multi();`, + Expected: []sql.Row{{"1/a"}}, + }, + }, + }, + }) +}