chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,293 @@
|
||||
// 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 node
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
"github.com/dolthub/go-mysql-server/sql/plan"
|
||||
|
||||
"github.com/dolthub/doltgresql/core/triggers"
|
||||
pgexprs "github.com/dolthub/doltgresql/server/expression"
|
||||
"github.com/dolthub/doltgresql/server/functions/framework"
|
||||
"github.com/dolthub/doltgresql/server/plpgsql"
|
||||
pgtypes "github.com/dolthub/doltgresql/server/types"
|
||||
)
|
||||
|
||||
// TriggerExecutionRowHandling states how to interpret the source row, or how to return the resulting row.
|
||||
type TriggerExecutionRowHandling uint8
|
||||
|
||||
const (
|
||||
TriggerExecutionRowHandling_None TriggerExecutionRowHandling = iota
|
||||
TriggerExecutionRowHandling_Old
|
||||
TriggerExecutionRowHandling_OldNew
|
||||
TriggerExecutionRowHandling_NewOld
|
||||
TriggerExecutionRowHandling_New
|
||||
)
|
||||
|
||||
// TriggerExecution handles the execution of a set of triggers on a table.
|
||||
type TriggerExecution struct {
|
||||
Timing triggers.TriggerTiming
|
||||
Triggers []triggers.Trigger
|
||||
Split TriggerExecutionRowHandling // How the source row should be split
|
||||
Return TriggerExecutionRowHandling // How the returned rows should be combined
|
||||
Sch sql.Schema
|
||||
Source sql.Node
|
||||
Runner pgexprs.StatementRunner
|
||||
}
|
||||
|
||||
var _ sql.ExecBuilderNode = (*TriggerExecution)(nil)
|
||||
var _ sql.Expressioner = (*TriggerExecution)(nil)
|
||||
|
||||
func (te *TriggerExecution) Children() []sql.Node {
|
||||
return []sql.Node{te.Source}
|
||||
}
|
||||
|
||||
// Expressions implements the interface sql.Expressioner.
|
||||
func (te *TriggerExecution) Expressions() []sql.Expression {
|
||||
return []sql.Expression{te.Runner}
|
||||
}
|
||||
|
||||
// IsReadOnly implements the interface sql.ExecBuilderNode.
|
||||
func (te *TriggerExecution) IsReadOnly() bool {
|
||||
return te.Source.IsReadOnly()
|
||||
}
|
||||
|
||||
// Resolved implements the interface sql.ExecBuilderNode.
|
||||
func (te *TriggerExecution) Resolved() bool {
|
||||
return te.Source.Resolved()
|
||||
}
|
||||
|
||||
// BuildRowIter implements the interface sql.ExecBuilderNode.
|
||||
func (te *TriggerExecution) BuildRowIter(ctx *sql.Context, b sql.NodeExecBuilder, r sql.Row) (sql.RowIter, error) {
|
||||
sourceIter, err := b.Build(ctx, te.Source, r)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// If there are no triggers, then we'll just return the source iter
|
||||
if len(te.Triggers) == 0 {
|
||||
return sourceIter, nil
|
||||
}
|
||||
trigFuncs := make([]framework.InterpretedFunction, len(te.Triggers))
|
||||
whens := make([]framework.InterpretedFunction, len(te.Triggers))
|
||||
for i, trig := range te.Triggers {
|
||||
trigFuncs[i], err = te.loadTriggerFunction(ctx, trig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// If we have a WHEN expression, then we need to build a "function" to execute the expression
|
||||
if len(trig.When) > 0 {
|
||||
whens[i] = framework.InterpretedFunction{
|
||||
ID: trigFuncs[i].ID, // Assign the same ID just so we have a valid one for later
|
||||
ReturnType: pgtypes.Bool,
|
||||
Statements: trig.When,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tgOp := ""
|
||||
switch te.Source.(type) {
|
||||
case *plan.InsertInto:
|
||||
tgOp = "INSERT"
|
||||
case *plan.Update:
|
||||
tgOp = "UPDATE"
|
||||
case *plan.DeleteFrom:
|
||||
tgOp = "DELETE"
|
||||
case *plan.Truncate:
|
||||
tgOp = "TRUNCATE"
|
||||
}
|
||||
|
||||
return &triggerExecutionIter{
|
||||
functions: trigFuncs,
|
||||
whens: whens,
|
||||
split: te.Split,
|
||||
treturn: te.Return,
|
||||
runner: te.Runner.Runner,
|
||||
sch: te.Sch,
|
||||
source: sourceIter,
|
||||
tgOp: tgOp,
|
||||
timing: te.Timing,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Schema implements the interface sql.ExecBuilderNode.
|
||||
func (te *TriggerExecution) Schema(ctx *sql.Context) sql.Schema {
|
||||
return te.Source.Schema(ctx)
|
||||
}
|
||||
|
||||
// String implements the interface sql.ExecBuilderNode.
|
||||
func (te *TriggerExecution) String() string {
|
||||
return "TRIGGER EXECUTION"
|
||||
}
|
||||
|
||||
// WithChildren implements the interface sql.ExecBuilderNode.
|
||||
func (te *TriggerExecution) WithChildren(ctx *sql.Context, children ...sql.Node) (sql.Node, error) {
|
||||
if len(children) != 1 {
|
||||
return nil, sql.ErrInvalidChildrenNumber.New(te, len(children), 1)
|
||||
}
|
||||
newTe := *te
|
||||
newTe.Source = children[0]
|
||||
return &newTe, nil
|
||||
}
|
||||
|
||||
// WithExpressions implements the interface sql.Expressioner.
|
||||
func (te *TriggerExecution) WithExpressions(ctx *sql.Context, expressions ...sql.Expression) (sql.Node, error) {
|
||||
if len(expressions) != 1 {
|
||||
return nil, sql.ErrInvalidChildrenNumber.New(te, len(expressions), 1)
|
||||
}
|
||||
newTe := *te
|
||||
newTe.Runner = expressions[0].(pgexprs.StatementRunner)
|
||||
return &newTe, nil
|
||||
}
|
||||
|
||||
// loadTriggerFunction loads the given trigger's framework.InterpretedFunction.
|
||||
func (te *TriggerExecution) loadTriggerFunction(ctx *sql.Context, trigger triggers.Trigger) (framework.InterpretedFunction, error) {
|
||||
function, err := loadFunction(ctx, nil, trigger.Function)
|
||||
if err != nil {
|
||||
return framework.InterpretedFunction{}, err
|
||||
}
|
||||
if !function.ID.IsValid() {
|
||||
return framework.InterpretedFunction{}, errors.Errorf("function %s() does not exist", trigger.Function.FunctionName())
|
||||
}
|
||||
if function.ReturnType != pgtypes.Trigger.ID {
|
||||
return framework.InterpretedFunction{}, errors.Errorf(`function %s must return type trigger`, function.ID.FunctionName())
|
||||
}
|
||||
return framework.InterpretedFunction{
|
||||
ID: function.ID,
|
||||
ReturnType: pgtypes.Trigger,
|
||||
ParameterNames: nil,
|
||||
ParameterTypes: nil,
|
||||
Variadic: function.Variadic,
|
||||
IsNonDeterministic: function.IsNonDeterministic,
|
||||
Strict: function.Strict,
|
||||
SRF: false,
|
||||
Statements: function.Operations,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// triggerExecutionIter is the iterator for TriggerExecution.
|
||||
type triggerExecutionIter struct {
|
||||
functions []framework.InterpretedFunction
|
||||
whens []framework.InterpretedFunction
|
||||
split TriggerExecutionRowHandling
|
||||
treturn TriggerExecutionRowHandling
|
||||
runner sql.StatementRunner
|
||||
sch sql.Schema
|
||||
source sql.RowIter
|
||||
tgOp string
|
||||
timing triggers.TriggerTiming
|
||||
}
|
||||
|
||||
var _ sql.RowIter = (*triggerExecutionIter)(nil)
|
||||
|
||||
// Next implements the interface sql.RowIter.
|
||||
func (t *triggerExecutionIter) Next(ctx *sql.Context) (sql.Row, error) {
|
||||
nextRow, err := t.source.Next(ctx)
|
||||
if err != nil {
|
||||
return nextRow, err
|
||||
}
|
||||
var oldRow sql.Row
|
||||
var newRow sql.Row
|
||||
switch t.split {
|
||||
case TriggerExecutionRowHandling_Old:
|
||||
oldRow = nextRow
|
||||
case TriggerExecutionRowHandling_OldNew:
|
||||
oldRow = nextRow[:len(t.sch)]
|
||||
newRow = nextRow[len(t.sch):]
|
||||
case TriggerExecutionRowHandling_NewOld:
|
||||
newRow = nextRow[:len(t.sch)]
|
||||
oldRow = nextRow[len(t.sch):]
|
||||
case TriggerExecutionRowHandling_New:
|
||||
newRow = nextRow
|
||||
}
|
||||
|
||||
// TODO: handle other special variables
|
||||
triggerVars := make(map[string]any)
|
||||
if t.tgOp != "" {
|
||||
triggerVars["TG_OP"] = t.tgOp
|
||||
}
|
||||
|
||||
for funcIdx, function := range t.functions {
|
||||
if t.whens[funcIdx].ID.IsValid() {
|
||||
whenValue, err := plpgsql.TriggerCall(ctx, t.whens[funcIdx], t.runner, t.sch, oldRow, newRow, triggerVars)
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "no valid cast for return value") {
|
||||
// TODO: this error should technically be caught during parsing, but interpreted functions don't
|
||||
// have the ability to determine types during parsing yet (also applies to the same error below)
|
||||
return nil, fmt.Errorf("argument of WHEN must be type boolean")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
whenBool, ok := whenValue.(bool)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("argument of WHEN must be type boolean")
|
||||
}
|
||||
if !whenBool {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
returnedValue, err := plpgsql.TriggerCall(ctx, function, t.runner, t.sch, oldRow, newRow, triggerVars)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if returnedValue == nil {
|
||||
// a returned value of NULL on a BEFORE trigger means to not modify the row, so we return a signal error
|
||||
if t.timing == triggers.TriggerTiming_Before {
|
||||
return nil, sql.ErrRowEditCanceled.New()
|
||||
} else {
|
||||
return nextRow, nil
|
||||
}
|
||||
}
|
||||
var ok bool
|
||||
returnedRow, ok := returnedValue.(sql.Row)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid trigger return value")
|
||||
}
|
||||
switch t.split {
|
||||
case TriggerExecutionRowHandling_Old:
|
||||
oldRow = returnedRow
|
||||
case TriggerExecutionRowHandling_OldNew, TriggerExecutionRowHandling_NewOld, TriggerExecutionRowHandling_New:
|
||||
newRow = returnedRow
|
||||
}
|
||||
}
|
||||
switch t.treturn {
|
||||
case TriggerExecutionRowHandling_Old:
|
||||
return oldRow, nil
|
||||
case TriggerExecutionRowHandling_OldNew:
|
||||
retRow := make(sql.Row, len(nextRow))
|
||||
copy(retRow, oldRow)
|
||||
copy(retRow[len(oldRow):], newRow)
|
||||
return retRow, nil
|
||||
case TriggerExecutionRowHandling_NewOld:
|
||||
retRow := make(sql.Row, len(nextRow))
|
||||
copy(retRow, newRow)
|
||||
copy(retRow[len(newRow):], oldRow)
|
||||
return retRow, nil
|
||||
case TriggerExecutionRowHandling_New:
|
||||
return newRow, nil
|
||||
default:
|
||||
return nextRow, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Close implements the interface sql.RowIter.
|
||||
func (t *triggerExecutionIter) Close(ctx *sql.Context) error {
|
||||
return t.source.Close(ctx)
|
||||
}
|
||||
Reference in New Issue
Block a user