chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,504 @@
|
||||
// Copyright 2025 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 plpgsql
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
"github.com/jackc/pgx/v5/pgproto3"
|
||||
|
||||
"github.com/dolthub/doltgresql/core/id"
|
||||
"github.com/dolthub/doltgresql/core/typecollection"
|
||||
"github.com/dolthub/doltgresql/postgres/parser/types"
|
||||
pgtypes "github.com/dolthub/doltgresql/server/types"
|
||||
)
|
||||
|
||||
// InterpretedFunction is an interface that essentially mirrors the implementation of InterpretedFunction in the
|
||||
// framework package.
|
||||
type InterpretedFunction interface {
|
||||
ApplyBindings(ctx *sql.Context, stack InterpreterStack, stmt string, bindings []string, enforceType bool) (newStmt string, varFound bool, err error)
|
||||
GetParameters() []*pgtypes.DoltgresType
|
||||
GetParameterNames() []string
|
||||
GetReturn() *pgtypes.DoltgresType
|
||||
GetStatements() []InterpreterOperation
|
||||
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
|
||||
// import cycles.
|
||||
var GetTypesCollectionFromContext func(ctx *sql.Context, database string) (*typecollection.TypeCollection, error)
|
||||
|
||||
// Call runs the contained operations on the given runner.
|
||||
func Call(ctx *sql.Context, iFunc InterpretedFunction, runner sql.StatementRunner, paramsAndReturn []*pgtypes.DoltgresType, vals []any) (any, error) {
|
||||
// Set up the initial state of the function
|
||||
stack := NewInterpreterStack(runner)
|
||||
// Add the parameters
|
||||
parameterTypes := iFunc.GetParameters()
|
||||
parameterNames := iFunc.GetParameterNames()
|
||||
if len(vals) != len(parameterTypes) {
|
||||
return nil, fmt.Errorf("parameter count mismatch: expected %d got %d", len(parameterTypes), len(vals))
|
||||
}
|
||||
for i := range vals {
|
||||
stack.NewVariableWithValue(parameterNames[i], parameterTypes[i], vals[i])
|
||||
}
|
||||
return call(ctx, iFunc, stack)
|
||||
}
|
||||
|
||||
// TriggerCall runs the contained trigger operations on the given runner.
|
||||
func TriggerCall(ctx *sql.Context, iFunc InterpretedFunction, runner sql.StatementRunner, sch sql.Schema, oldRow sql.Row, newRow sql.Row, trigVars map[string]any) (any, error) {
|
||||
// Set up the initial state of the function
|
||||
stack := NewInterpreterStack(runner)
|
||||
// Add the special variables
|
||||
stack.NewRecord("OLD", sch, oldRow)
|
||||
stack.NewRecord("NEW", sch, newRow)
|
||||
for varName, val := range trigVars {
|
||||
varType, ok := triggerSpecialVariables[varName]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown variable %s for trigger", varName)
|
||||
}
|
||||
stack.NewVariableWithValue(varName, varType, val)
|
||||
}
|
||||
return call(ctx, iFunc, stack)
|
||||
}
|
||||
|
||||
// call runs the contained operations on the given runner.
|
||||
func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (any, error) {
|
||||
// We increment before accessing, so start at -1
|
||||
counter := -1
|
||||
// Run the statements
|
||||
statements := iFunc.GetStatements()
|
||||
for {
|
||||
counter++
|
||||
if counter >= len(statements) {
|
||||
break
|
||||
} else if counter < 0 {
|
||||
panic("negative function counter")
|
||||
}
|
||||
|
||||
operation := statements[counter]
|
||||
switch operation.OpCode {
|
||||
case OpCode_Alias:
|
||||
iv := stack.GetVariable(operation.PrimaryData)
|
||||
if iv.Type == nil {
|
||||
return nil, fmt.Errorf("variable `%s` could not be found", operation.PrimaryData)
|
||||
}
|
||||
stack.NewVariableAlias(operation.Target, operation.PrimaryData)
|
||||
case OpCode_Assign:
|
||||
iv := stack.GetVariable(operation.Target)
|
||||
if iv.Type == nil {
|
||||
return nil, fmt.Errorf("variable `%s` could not be found", operation.Target)
|
||||
}
|
||||
retVal, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, iv.Type, operation.SecondaryData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
err = stack.SetVariable(ctx, operation.Target, retVal)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
case OpCode_Declare:
|
||||
typeCollection, err := GetTypesCollectionFromContext(ctx, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// pg_query_go sets PrimaryData for implicit CASE statement variables to
|
||||
// `pg_catalog."integer"`, so we remove double-quotes and extract the schema name.
|
||||
typeName := operation.PrimaryData
|
||||
typeName = strings.ReplaceAll(typeName, `"`, "")
|
||||
schemaName := "pg_catalog"
|
||||
if strings.Contains(typeName, ".") {
|
||||
parts := strings.Split(typeName, ".")
|
||||
schemaName = parts[0]
|
||||
typeName = parts[1]
|
||||
// Check the NonKeyword type names to see if we're looking at
|
||||
// an alias of a type if we're in the pg_catalog schema.
|
||||
// Skip array types (names starting with "_") since their internal
|
||||
// lookup key uses the "_typename" form, not the "typename[]" form
|
||||
// that TypeForNonKeywordTypeName returns.
|
||||
if schemaName == "pg_catalog" && !strings.HasPrefix(typeName, "_") {
|
||||
typ, ok, _ := types.TypeForNonKeywordTypeName(typeName)
|
||||
if ok && typ != nil {
|
||||
typeName = typ.Name()
|
||||
}
|
||||
}
|
||||
}
|
||||
resolvedType, err := typeCollection.GetType(ctx, id.NewType(schemaName, typeName))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resolvedType == nil {
|
||||
return nil, pgtypes.ErrTypeDoesNotExist.New(operation.PrimaryData)
|
||||
}
|
||||
if len(operation.SecondaryData) != 0 {
|
||||
defVal := operation.SecondaryData[0]
|
||||
// Default value can be a literal value or a reference to parameter
|
||||
isParam := false
|
||||
for _, param := range iFunc.GetParameterNames() {
|
||||
if param == defVal {
|
||||
isParam = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if isParam {
|
||||
ivr := stack.GetVariable(defVal)
|
||||
if ivr.Value != nil {
|
||||
stack.NewVariableWithValue(operation.Target, resolvedType, *ivr.Value)
|
||||
} else {
|
||||
stack.NewVariable(operation.Target, resolvedType)
|
||||
}
|
||||
} else {
|
||||
val, err := resolvedType.IoInput(ctx, strings.Trim(operation.SecondaryData[0], "'"))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stack.NewVariableWithValue(operation.Target, resolvedType, val)
|
||||
}
|
||||
} else {
|
||||
stack.NewVariable(operation.Target, resolvedType)
|
||||
}
|
||||
case OpCode_DeleteInto:
|
||||
// TODO: implement
|
||||
case OpCode_Exception:
|
||||
// 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)
|
||||
if target.Type == nil {
|
||||
return nil, fmt.Errorf("variable `%s` could not be found", operation.Target)
|
||||
}
|
||||
retVal, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, target.Type, operation.SecondaryData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
err = stack.SetVariable(ctx, operation.Target, retVal)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
} else {
|
||||
_, _, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
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 {
|
||||
for ; counter < operation.Index-1; counter++ {
|
||||
switch statements[counter].OpCode {
|
||||
case OpCode_ScopeBegin:
|
||||
stack.PushScope()
|
||||
case OpCode_ScopeEnd:
|
||||
stack.PopScope()
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for ; counter > operation.Index-1; counter-- {
|
||||
switch statements[counter].OpCode {
|
||||
case OpCode_ScopeBegin:
|
||||
stack.PopScope()
|
||||
case OpCode_ScopeEnd:
|
||||
stack.PushScope()
|
||||
}
|
||||
}
|
||||
}
|
||||
case OpCode_If:
|
||||
retVal, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, pgtypes.Bool, operation.SecondaryData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if retVal.(bool) {
|
||||
// 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)
|
||||
if 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.
|
||||
// https://www.postgresql.org/docs/current/runtime-config-client.html#GUC-CLIENT-MIN-MESSAGES
|
||||
|
||||
message, err := evaluteNoticeMessage(ctx, iFunc, operation, stack)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if operation.PrimaryData == "EXCEPTION" {
|
||||
// TODO: Notices at the EXCEPTION level should also abort the current tx.
|
||||
return nil, errors.New(message)
|
||||
} else {
|
||||
noticeResponse := &pgproto3.NoticeResponse{
|
||||
Severity: operation.PrimaryData,
|
||||
Message: message,
|
||||
}
|
||||
if err = applyNoticeOptions(ctx, noticeResponse, operation.Options); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sess := dsess.DSessFromSess(ctx.Session)
|
||||
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), " ", "")
|
||||
if normalized == "select$1;" {
|
||||
if strings.EqualFold(operation.SecondaryData[0], "new") {
|
||||
return *stack.GetVariable("NEW").Value, nil
|
||||
} else if strings.EqualFold(operation.SecondaryData[0], "old") {
|
||||
return *stack.GetVariable("OLD").Value, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
val, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, iFunc.GetReturn(), operation.SecondaryData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 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_ForQueryInit:
|
||||
schema, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stack.InitCursor(operation.Target, schema, rows)
|
||||
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.
|
||||
counter = operation.Index - 1
|
||||
} else {
|
||||
if err := stack.UpdateRecord(operation.Target, schema, row); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
case OpCode_ReturnQuery:
|
||||
schema, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
records, err := convertRowsToRecords(schema, rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stack.BufferReturnQueryResults(records)
|
||||
case OpCode_ScopeBegin:
|
||||
stack.PushScope()
|
||||
case OpCode_ScopeEnd:
|
||||
stack.PopScope()
|
||||
case OpCode_SelectInto:
|
||||
// TODO: implement
|
||||
case OpCode_UpdateInto:
|
||||
// TODO: implement
|
||||
default:
|
||||
panic("unimplemented opcode")
|
||||
}
|
||||
}
|
||||
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, error) {
|
||||
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 {
|
||||
// non-Doltgres types are still used in analysis, but we only support disk serialization
|
||||
// for Doltgres types, so we must convert the GMS type to the nearest Doltgres type here.
|
||||
// TODO: this conversion isn't fully accurate. expression.GMSCast has additional logic in
|
||||
// its Eval() method to handle types more exactly and also handles converting the
|
||||
// value to ensure it is well formed for the returned DoltgresType. We can't
|
||||
// currently use GMSCast directly here though, because of a dependency cycle, so
|
||||
// that conversion logic needs to be extracted into a package both places can import.
|
||||
var err error
|
||||
doltgresType, err = pgtypes.FromGmsTypeToDoltgresType(t)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
record[i] = pgtypes.RecordValue{
|
||||
Value: field,
|
||||
Type: doltgresType,
|
||||
}
|
||||
}
|
||||
records = append(records, record)
|
||||
}
|
||||
|
||||
return records, nil
|
||||
}
|
||||
|
||||
// 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 {
|
||||
i, err := strconv.Atoi(key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
switch NoticeOptionType(i) {
|
||||
case NoticeOptionTypeErrCode:
|
||||
noticeResponse.Code = value
|
||||
case NoticeOptionTypeMessage:
|
||||
noticeResponse.Message = value
|
||||
case NoticeOptionTypeDetail:
|
||||
noticeResponse.Detail = value
|
||||
case NoticeOptionTypeHint:
|
||||
noticeResponse.Hint = value
|
||||
case NoticeOptionTypeConstraint:
|
||||
noticeResponse.ConstraintName = value
|
||||
case NoticeOptionTypeDataType:
|
||||
noticeResponse.DataTypeName = value
|
||||
case NoticeOptionTypeTable:
|
||||
noticeResponse.TableName = value
|
||||
case NoticeOptionTypeSchema:
|
||||
noticeResponse.SchemaName = value
|
||||
default:
|
||||
ctx.GetLogger().Warnf("unhandled notice option type: %s", key)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// evaluteNoticeMessage evaluates the message for a RAISE NOTICE statement, including
|
||||
// evaluating any specified parameters and plugging them into the message in place of
|
||||
// the % placeholders.
|
||||
func evaluteNoticeMessage(ctx *sql.Context, iFunc InterpretedFunction,
|
||||
operation InterpreterOperation, stack InterpreterStack) (string, error) {
|
||||
message := operation.SecondaryData[0]
|
||||
if len(operation.SecondaryData) > 1 {
|
||||
params := operation.SecondaryData[1:]
|
||||
currentParamIdx := 0
|
||||
|
||||
parts := strings.Split(message, "%%")
|
||||
for i, part := range parts {
|
||||
for strings.Contains(part, "%") {
|
||||
if currentParamIdx >= len(params) {
|
||||
return "", errors.New("too few parameters specified for RAISE")
|
||||
}
|
||||
currentParam := params[currentParamIdx]
|
||||
currentParamIdx += 1
|
||||
formattedVar, varFound, err := iFunc.ApplyBindings(ctx, stack, "$1", []string{currentParam}, false)
|
||||
if varFound {
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
part = strings.Replace(part, "%", formattedVar, 1)
|
||||
} else {
|
||||
retVal, err := iFunc.QuerySingleReturn(ctx, stack, fmt.Sprintf("SELECT (%s)::text", currentParam), nil, nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
stringVal := fmt.Sprintf("%v", retVal) // We should always return a string, but this is just a safety net
|
||||
part = strings.Replace(part, "%", stringVal, 1)
|
||||
}
|
||||
}
|
||||
parts[i] = part
|
||||
}
|
||||
if currentParamIdx < len(params) {
|
||||
return "", errors.New("too many parameters specified for RAISE")
|
||||
}
|
||||
message = strings.Join(parts, "%")
|
||||
}
|
||||
return message, nil
|
||||
}
|
||||
|
||||
// triggerSpecialVariables are the list of special variables for triggers.
|
||||
// https://www.postgresql.org/docs/15/plpgsql-trigger.html
|
||||
// TODO: NEW and OLD variables are handled separately using `InterpreterStack.NewRecord` function.
|
||||
var triggerSpecialVariables = map[string]*pgtypes.DoltgresType{
|
||||
//"NEW":
|
||||
//"OLD":
|
||||
"TG_NAME": pgtypes.Name,
|
||||
"TG_WHEN": pgtypes.Text,
|
||||
"TG_LEVEL": pgtypes.Text,
|
||||
"TG_OP": pgtypes.Text,
|
||||
"TG_RELID": pgtypes.Oid,
|
||||
"TG_RELNAME": pgtypes.Name,
|
||||
"TG_TABLE_NAME": pgtypes.Name,
|
||||
"TG_TABLE_SCHEMA": pgtypes.Name,
|
||||
"TG_NARGS": pgtypes.Int32,
|
||||
"TG_ARGV[]": pgtypes.TextArray,
|
||||
}
|
||||
Reference in New Issue
Block a user