chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,300 @@
|
||||
// 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"
|
||||
"strings"
|
||||
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
|
||||
pgtypes "github.com/dolthub/doltgresql/server/types"
|
||||
"github.com/dolthub/doltgresql/utils"
|
||||
)
|
||||
|
||||
// cursorState holds the result set for a FOR record IN query LOOP cursor.
|
||||
type cursorState struct {
|
||||
Schema sql.Schema
|
||||
Rows []sql.Row
|
||||
Index int
|
||||
}
|
||||
|
||||
// interpreterVariable is a variable that lives on the stack. This will hold an actual value, but will not be directly
|
||||
// interacted with. InterpreterVariableReference are, instead, the avenue of interaction as a variable may be an
|
||||
// aggregate type (such as a record).
|
||||
type interpreterVariable struct {
|
||||
Record sql.Schema // TODO: all records carry their type information alongside the value, so this is redundant
|
||||
Type *pgtypes.DoltgresType
|
||||
Value any
|
||||
}
|
||||
|
||||
// InterpreterVariableReference is a reference to a variable that lives on the stack. If the type is not null, then it
|
||||
// is valid to dereference the value for assignment. We make use of references rather than directly interacting with
|
||||
// the variables as this allows for interacting with sections of aggregate types (such as record) as well as normal
|
||||
// variable interaction.
|
||||
type InterpreterVariableReference struct {
|
||||
Type *pgtypes.DoltgresType
|
||||
Value *any
|
||||
}
|
||||
|
||||
// InterpreterScopeDetails contains all of the details that are relevant to a particular scope.
|
||||
type InterpreterScopeDetails struct {
|
||||
variables map[string]*interpreterVariable
|
||||
label string
|
||||
}
|
||||
|
||||
// InterpreterStack represents the working information that an interpreter will use during execution. It is not exactly
|
||||
// the same as a stack in the traditional programming sense, but rather is a loose abstraction that serves the same
|
||||
// general purpose.
|
||||
type InterpreterStack struct {
|
||||
stack *utils.Stack[*InterpreterScopeDetails]
|
||||
runner sql.StatementRunner
|
||||
labelID int
|
||||
|
||||
// returnQueryBuffer buffers results from RETURN QUERY statements
|
||||
returnQueryBuffer [][]pgtypes.RecordValue
|
||||
// cursors holds the active FOR record IN query LOOP result sets
|
||||
cursors map[string]*cursorState
|
||||
}
|
||||
|
||||
// NewInterpreterStack creates a new InterpreterStack.
|
||||
func NewInterpreterStack(runner sql.StatementRunner) InterpreterStack {
|
||||
stack := utils.NewStack[*InterpreterScopeDetails]()
|
||||
// This first push represents the function base, including parameters
|
||||
stack.Push(&InterpreterScopeDetails{
|
||||
variables: make(map[string]*interpreterVariable),
|
||||
})
|
||||
return InterpreterStack{
|
||||
stack: stack,
|
||||
runner: runner,
|
||||
cursors: make(map[string]*cursorState),
|
||||
}
|
||||
}
|
||||
|
||||
// Details returns the details for the current scope.
|
||||
func (is *InterpreterStack) Details() *InterpreterScopeDetails {
|
||||
return is.stack.Peek()
|
||||
}
|
||||
|
||||
// Runner returns the runner that is being used for the function's execution.
|
||||
func (is *InterpreterStack) Runner() sql.StatementRunner {
|
||||
return is.runner
|
||||
}
|
||||
|
||||
// GetCurrentLabel traverses the stack (starting from the top) returning the first label found. Returns an empty string
|
||||
// if no labels were set.
|
||||
func (is *InterpreterStack) GetCurrentLabel() string {
|
||||
for i := 0; i < is.stack.Len(); i++ {
|
||||
label := is.stack.PeekDepth(i).label
|
||||
if len(label) > 0 {
|
||||
return label
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// GetVariable traverses the stack (starting from the top) to find a variable with a matching name. Returns nil if no
|
||||
// variable was found.
|
||||
func (is *InterpreterStack) GetVariable(name string) InterpreterVariableReference {
|
||||
// TODO: handle nested record access
|
||||
fieldName := ""
|
||||
if strings.Count(name, ".") == 1 {
|
||||
splitName := strings.Split(name, ".")
|
||||
name = splitName[0]
|
||||
fieldName = splitName[1]
|
||||
}
|
||||
for i := 0; i < is.stack.Len(); i++ {
|
||||
if iv, ok := is.stack.PeekDepth(i).variables[name]; ok {
|
||||
if len(fieldName) == 0 {
|
||||
return InterpreterVariableReference{
|
||||
Type: iv.Type,
|
||||
Value: &iv.Value,
|
||||
}
|
||||
} else if len(iv.Record) > 0 {
|
||||
fieldIdx := iv.Record.IndexOf(fieldName, iv.Record[0].Source)
|
||||
if fieldIdx == -1 {
|
||||
// TODO: implement this as a proper error for missing record field rather than the generic "variable not found"
|
||||
return InterpreterVariableReference{}
|
||||
}
|
||||
return InterpreterVariableReference{
|
||||
Type: iv.Record[fieldIdx].Type.(*pgtypes.DoltgresType),
|
||||
Value: &(iv.Value.(sql.Row)[fieldIdx]),
|
||||
}
|
||||
} else if iv.Type.IsCompositeType() {
|
||||
for fieldIdx := range iv.Type.CompositeAttrs {
|
||||
if iv.Type.CompositeAttrs[fieldIdx].Name == fieldName {
|
||||
vals := iv.Value.([]pgtypes.RecordValue)
|
||||
return InterpreterVariableReference{
|
||||
Type: vals[fieldIdx].Type.(*pgtypes.DoltgresType),
|
||||
Value: &(vals[fieldIdx].Value),
|
||||
}
|
||||
}
|
||||
}
|
||||
// The field could not be found
|
||||
return InterpreterVariableReference{}
|
||||
} else {
|
||||
// Can't access fields on an empty record
|
||||
return InterpreterVariableReference{}
|
||||
}
|
||||
}
|
||||
}
|
||||
return InterpreterVariableReference{}
|
||||
}
|
||||
|
||||
// ListVariables returns a map with the names of all variables. The attached slice represents field names for records.
|
||||
// All names are lowercased.
|
||||
func (is *InterpreterStack) ListVariables() map[string][]string {
|
||||
seen := make(map[string][]string)
|
||||
for i := 0; i < is.stack.Len(); i++ {
|
||||
for varName, iv := range is.stack.PeekDepth(i).variables {
|
||||
var fieldNames []string
|
||||
if len(iv.Record) > 0 {
|
||||
for _, col := range iv.Record {
|
||||
fieldNames = append(fieldNames, strings.ToLower(col.Name))
|
||||
}
|
||||
}
|
||||
seen[strings.ToLower(varName)] = fieldNames
|
||||
}
|
||||
}
|
||||
return seen
|
||||
}
|
||||
|
||||
// NewRecord creates a new record in the current scope. If a record with the same name exists in a previous scope, then
|
||||
// that record will be shadowed until the current scope exits.
|
||||
func (is *InterpreterStack) NewRecord(name string, sch sql.Schema, val sql.Row) {
|
||||
// TODO: this is currently implemented only for the specific record types used in triggers: OLD and NEW
|
||||
var newVal sql.Row
|
||||
if val != nil {
|
||||
newVal = make(sql.Row, len(val))
|
||||
copy(newVal, val)
|
||||
}
|
||||
is.stack.Peek().variables[name] = &interpreterVariable{
|
||||
Record: sch,
|
||||
Type: pgtypes.Trigger, // TODO: we need to implement the RECORD pseudotype and replace the TRIGGER type here
|
||||
Value: newVal,
|
||||
}
|
||||
}
|
||||
|
||||
// NewVariable creates a new variable in the current scope. If a variable with the same name exists in a previous scope,
|
||||
// then that variable will be shadowed until the current scope exits.
|
||||
func (is *InterpreterStack) NewVariable(name string, typ *pgtypes.DoltgresType) {
|
||||
is.NewVariableWithValue(name, typ, typ.Zero())
|
||||
}
|
||||
|
||||
// NewVariableWithValue creates a new variable in the current scope, setting its initial value to the one given.
|
||||
func (is *InterpreterStack) NewVariableWithValue(name string, typ *pgtypes.DoltgresType, val any) {
|
||||
is.stack.Peek().variables[name] = &interpreterVariable{
|
||||
Type: typ,
|
||||
Value: val,
|
||||
}
|
||||
}
|
||||
|
||||
// NewVariableAlias creates a new variable alias, named |alias|, in the current frame of this stack,
|
||||
// pointing to the specified |variable|.
|
||||
func (is *InterpreterStack) NewVariableAlias(alias string, target string) {
|
||||
for i := 0; i < is.stack.Len(); i++ {
|
||||
if iv, ok := is.stack.PeekDepth(i).variables[target]; ok {
|
||||
// TODO: this won't work for RECORD types
|
||||
is.stack.Peek().variables[alias] = iv
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// PushScope creates a new scope.
|
||||
func (is *InterpreterStack) PushScope() {
|
||||
is.stack.Push(&InterpreterScopeDetails{
|
||||
variables: make(map[string]*interpreterVariable),
|
||||
})
|
||||
}
|
||||
|
||||
// PopScope removes the current scope.
|
||||
func (is *InterpreterStack) PopScope() {
|
||||
is.stack.Pop()
|
||||
}
|
||||
|
||||
// SetVariable sets the first variable found, with a matching name, to the value given. This does not ensure that the
|
||||
// value matches the expectations of the type, so it should be validated before this is called. Returns an error if the
|
||||
// variable cannot be found.
|
||||
func (is *InterpreterStack) SetVariable(ctx *sql.Context, name string, val any) error {
|
||||
iv := is.GetVariable(name)
|
||||
if iv.Type == nil {
|
||||
return fmt.Errorf("variable `%s` could not be found", name)
|
||||
}
|
||||
*iv.Value = val
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetLabel sets the label for the current scope.
|
||||
func (is *InterpreterStack) SetLabel(label string) {
|
||||
is.stack.Peek().label = label
|
||||
}
|
||||
|
||||
// SetAnonymousLabel sets the label for the current scope to a guaranteed unique value.
|
||||
func (is *InterpreterStack) SetAnonymousLabel() {
|
||||
// Postgres labels cannot have a tab character, so we can generate a label with one to guarantee it's unique
|
||||
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
|
||||
}
|
||||
|
||||
// InitCursor stores the result set for a FOR record IN query LOOP cursor.
|
||||
func (is *InterpreterStack) InitCursor(name string, schema sql.Schema, rows []sql.Row) {
|
||||
is.cursors[name] = &cursorState{
|
||||
Schema: schema,
|
||||
Rows: rows,
|
||||
Index: 0,
|
||||
}
|
||||
}
|
||||
|
||||
// AdvanceCursor returns the next row for the named cursor and advances its index.
|
||||
// Returns (schema, row, true) if a row is available, or (nil, nil, false) when exhausted.
|
||||
func (is *InterpreterStack) AdvanceCursor(name string) (sql.Schema, sql.Row, bool) {
|
||||
cs, ok := is.cursors[name]
|
||||
if !ok || cs.Index >= len(cs.Rows) {
|
||||
return nil, nil, false
|
||||
}
|
||||
row := cs.Rows[cs.Index]
|
||||
cs.Index++
|
||||
return cs.Schema, row, true
|
||||
}
|
||||
|
||||
// CloseCursor removes the named cursor from the stack.
|
||||
func (is *InterpreterStack) CloseCursor(name string) {
|
||||
delete(is.cursors, name)
|
||||
}
|
||||
|
||||
// UpdateRecord finds the named variable and sets its schema and row value.
|
||||
func (is *InterpreterStack) UpdateRecord(name string, schema sql.Schema, val sql.Row) error {
|
||||
for i := 0; i < is.stack.Len(); i++ {
|
||||
if iv, ok := is.stack.PeekDepth(i).variables[name]; ok {
|
||||
iv.Record = schema
|
||||
iv.Value = val
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("record variable `%s` could not be found", name)
|
||||
}
|
||||
Reference in New Issue
Block a user