From 1a887995816df8225a1e1d6199e04e33f6559fc6 Mon Sep 17 00:00:00 2001 From: Aaron Son Date: Mon, 31 Aug 2026 13:30:01 +0200 Subject: [PATCH] Fix INTO semantics in plpgsql An INTO clause without STRICT has three behaviors that Doltgres got wrong. A query matching no rows raised instead of setting every target to NULL. A query matching several rows raised instead of keeping the first and discarding the rest. And a target had to have exactly the type the query column produced, rather than accepting anything with an assignment cast to it. An INTO naming several variables also assigned only the first of them, because the loop that walked the targets was indexed by row rather than by target. QueryRowReturn replaces the two paths that INTO used to take, so a single target and several targets are now handled the same way, and reports whether the query produced a row at all. The casting that QuerySingleReturn did inline is extracted as castQueryValue so both share it. --- .../framework/interpreted_function.go | 103 ++++++++++----- server/plpgsql/interpreter_logic.go | 52 +++----- testing/go/plpgsql_into_test.go | 124 ++++++++++++++++++ 3 files changed, 213 insertions(+), 66 deletions(-) create mode 100644 testing/go/plpgsql_into_test.go 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"}}, + }, + }, + }, + }) +}