chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,505 @@
|
||||
// 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"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
pg_query "github.com/dolthub/pg_query_go/v6"
|
||||
)
|
||||
|
||||
// Statement represents a PL/pgSQL statement.
|
||||
type Statement interface {
|
||||
// OperationSize reports the number of operations that the statement will convert to.
|
||||
OperationSize() int32
|
||||
// AppendOperations adds the statement to the operation slice.
|
||||
AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error
|
||||
}
|
||||
|
||||
// Assignment represents an assignment statement.
|
||||
type Assignment struct {
|
||||
VariableName string
|
||||
Expression string
|
||||
VariableIndex int32 // TODO: figure out what this is used for, probably to get around shadowed variables?
|
||||
}
|
||||
|
||||
var _ Statement = Assignment{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (Assignment) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt Assignment) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
expression, referencedVariables, err := substituteVariableReferences(stmt.Expression, stack)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Assign,
|
||||
PrimaryData: "SELECT " + expression + ";",
|
||||
SecondaryData: referencedVariables,
|
||||
Target: stmt.VariableName,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Block contains a collection of statements, alongside the variables that were declared for the block. Only the
|
||||
// top-level block will contain parameter variables.
|
||||
type Block struct {
|
||||
TriggerNew int32 // When non-zero, indicates that the NEW record exists for use with triggers
|
||||
TriggerOld int32 // When non-zero, indicates that the OLD record exists for use with triggers
|
||||
Variables []Variable
|
||||
Records []Record
|
||||
Body []Statement
|
||||
Label string
|
||||
IsLoop bool
|
||||
}
|
||||
|
||||
var _ Statement = Block{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (stmt Block) OperationSize() int32 {
|
||||
total := int32(2) // We start with 2 since we'll have ScopeBegin and ScopeEnd
|
||||
for _, variable := range stmt.Variables {
|
||||
if !variable.IsParameter {
|
||||
total++
|
||||
}
|
||||
}
|
||||
for _, innerStmt := range stmt.Body {
|
||||
total += innerStmt.OperationSize()
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt Block) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
stack.PushScope()
|
||||
stack.SetLabel(stmt.Label) // If the label is empty, then this won't change anything
|
||||
var loop string
|
||||
if stmt.IsLoop {
|
||||
loop = "_"
|
||||
// All loops need a label, so we'll make an anonymous one if an explicit one hasn't been given
|
||||
if len(stmt.Label) == 0 {
|
||||
stack.SetAnonymousLabel()
|
||||
stmt.Label = stack.GetCurrentLabel()
|
||||
}
|
||||
}
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_ScopeBegin,
|
||||
PrimaryData: stmt.Label,
|
||||
Target: loop,
|
||||
})
|
||||
for _, variable := range stmt.Variables {
|
||||
op := InterpreterOperation{
|
||||
OpCode: OpCode_Declare,
|
||||
PrimaryData: variable.Type,
|
||||
Target: variable.Name,
|
||||
}
|
||||
var val any
|
||||
if variable.Default != "" {
|
||||
op.SecondaryData = []string{variable.Default}
|
||||
val = variable.Default
|
||||
}
|
||||
if !variable.IsParameter {
|
||||
*ops = append(*ops, op)
|
||||
}
|
||||
stack.NewVariableWithValue(variable.Name, nil, val)
|
||||
}
|
||||
for _, record := range stmt.Records {
|
||||
var fakeSch sql.Schema
|
||||
for _, fieldName := range record.Fields {
|
||||
fakeSch = append(fakeSch, &sql.Column{Name: fieldName})
|
||||
}
|
||||
stack.NewRecord(record.Name, fakeSch, nil)
|
||||
}
|
||||
for _, innerStmt := range stmt.Body {
|
||||
if err := innerStmt.AppendOperations(ops, stack); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_ScopeEnd,
|
||||
})
|
||||
stack.PopScope()
|
||||
return nil
|
||||
}
|
||||
|
||||
// ExecuteSQL represents a standard SQL statement's execution (including the INTO syntax).
|
||||
type ExecuteSQL struct {
|
||||
Statement string
|
||||
Target string
|
||||
}
|
||||
|
||||
var _ Statement = ExecuteSQL{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (ExecuteSQL) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt ExecuteSQL) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
statementStr, referencedVariables, err := substituteVariableReferences(stmt.Statement, stack)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Execute,
|
||||
PrimaryData: statementStr,
|
||||
SecondaryData: referencedVariables,
|
||||
Target: stmt.Target,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// DynamicExecute represents a dynamic SQL statement's execution.
|
||||
type DynamicExecute struct {
|
||||
Query string
|
||||
Params []string
|
||||
Target string
|
||||
}
|
||||
|
||||
var _ Statement = DynamicExecute{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (DynamicExecute) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt DynamicExecute) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Execute,
|
||||
PrimaryData: stmt.Query,
|
||||
SecondaryData: stmt.Params,
|
||||
Target: stmt.Target,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// ForQueryInit executes a SQL query and stores the result set in a named cursor on the stack.
|
||||
// It is the first operation emitted for a FOR record IN query LOOP statement.
|
||||
type ForQueryInit struct {
|
||||
CursorName string
|
||||
Query string
|
||||
}
|
||||
|
||||
var _ Statement = ForQueryInit{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (ForQueryInit) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt ForQueryInit) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
queryStr, referencedVariables, err := substituteVariableReferences(stmt.Query, stack)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_ForQueryInit,
|
||||
PrimaryData: queryStr,
|
||||
SecondaryData: referencedVariables,
|
||||
Target: stmt.CursorName,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// ForQueryNext fetches the next row from a named cursor and assigns it to a record variable.
|
||||
// When the cursor is exhausted it jumps forward by GotoOffset (like an If), exiting the loop.
|
||||
type ForQueryNext struct {
|
||||
CursorName string
|
||||
RecordVar string
|
||||
GotoOffset int32
|
||||
}
|
||||
|
||||
var _ Statement = ForQueryNext{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (ForQueryNext) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt ForQueryNext) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_ForQueryNext,
|
||||
PrimaryData: stmt.CursorName,
|
||||
Target: stmt.RecordVar,
|
||||
Index: len(*ops) + int(stmt.GotoOffset),
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Goto jumps to the counter at the given offset.
|
||||
type Goto struct {
|
||||
Offset int32
|
||||
Label string
|
||||
NearestScopeOp bool
|
||||
}
|
||||
|
||||
var _ Statement = Goto{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (Goto) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt Goto) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
if len(stmt.Label) > 0 {
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Goto,
|
||||
PrimaryData: stmt.Label,
|
||||
Index: int(stmt.Offset),
|
||||
})
|
||||
} else if stmt.NearestScopeOp {
|
||||
label := stack.GetCurrentLabel()
|
||||
if len(label) == 0 {
|
||||
if stmt.Offset > 0 {
|
||||
return errors.New("EXIT cannot be used outside a loop, unless it has a label")
|
||||
} else {
|
||||
return errors.New("CONTINUE cannot be used outside a loop")
|
||||
}
|
||||
}
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Goto,
|
||||
PrimaryData: label,
|
||||
Index: int(stmt.Offset),
|
||||
})
|
||||
} else {
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Goto,
|
||||
Index: len(*ops) + int(stmt.Offset),
|
||||
})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// If represents an IF condition, alongside its Goto offset if the condition is true.
|
||||
type If struct {
|
||||
Condition string
|
||||
GotoOffset int32
|
||||
}
|
||||
|
||||
var _ Statement = If{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (If) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt If) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
condition, referencedVariables, err := substituteVariableReferences(stmt.Condition, stack)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_If,
|
||||
PrimaryData: "SELECT " + condition + ";",
|
||||
SecondaryData: referencedVariables,
|
||||
Index: len(*ops) + int(stmt.GotoOffset),
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Perform represents a PERFORM statement.
|
||||
type Perform struct {
|
||||
Statement string
|
||||
}
|
||||
|
||||
var _ Statement = Perform{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (Perform) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt Perform) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
statementStr, referencedVariables, err := substituteVariableReferences(stmt.Statement, stack)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Perform,
|
||||
PrimaryData: statementStr,
|
||||
SecondaryData: referencedVariables,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Raise represents a RAISE statement
|
||||
type Raise struct {
|
||||
Level string
|
||||
Message string
|
||||
Params []string
|
||||
Options map[string]string
|
||||
}
|
||||
|
||||
var _ Statement = Raise{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (r Raise) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (r Raise) AppendOperations(ops *[]InterpreterOperation, _ *InterpreterStack) error {
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Raise,
|
||||
PrimaryData: r.Level,
|
||||
SecondaryData: append([]string{r.Message}, r.Params...),
|
||||
Options: r.Options,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Record represents a record (along with known fields for future access). These are exclusively found within Block.
|
||||
type Record struct {
|
||||
Name string
|
||||
Fields []string
|
||||
}
|
||||
|
||||
// ReturnQuery represents a RETURN QUERY statement.
|
||||
type ReturnQuery struct {
|
||||
Query string
|
||||
}
|
||||
|
||||
var _ Statement = ReturnQuery{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (r ReturnQuery) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (r ReturnQuery) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
query, referencedVariables, err := substituteVariableReferences(r.Query, stack)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_ReturnQuery,
|
||||
PrimaryData: query,
|
||||
SecondaryData: referencedVariables,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Return represents a RETURN statement.
|
||||
type Return struct {
|
||||
Expression string
|
||||
}
|
||||
|
||||
var _ Statement = Return{}
|
||||
|
||||
// OperationSize implements the interface Statement.
|
||||
func (Return) OperationSize() int32 {
|
||||
return 1
|
||||
}
|
||||
|
||||
// AppendOperations implements the interface Statement.
|
||||
func (stmt Return) AppendOperations(ops *[]InterpreterOperation, stack *InterpreterStack) error {
|
||||
expression, referencedVariables, err := substituteVariableReferences(stmt.Expression, stack)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(expression) > 0 {
|
||||
expression = "SELECT " + expression + ";"
|
||||
}
|
||||
*ops = append(*ops, InterpreterOperation{
|
||||
OpCode: OpCode_Return,
|
||||
PrimaryData: expression,
|
||||
SecondaryData: referencedVariables,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Variable represents a variable. These are exclusively found within Block.
|
||||
type Variable struct {
|
||||
Name string
|
||||
Type string
|
||||
IsParameter bool
|
||||
Default string
|
||||
}
|
||||
|
||||
// OperationSizeForStatements returns the sum of OperationSize for every statement.
|
||||
func OperationSizeForStatements(stmts []Statement) int32 {
|
||||
total := int32(0)
|
||||
for _, stmt := range stmts {
|
||||
total += stmt.OperationSize()
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// substituteVariableReferences parses the specified |expression| and replaces
|
||||
// any token that matches a variable name in the |stack| with "$N", where N
|
||||
// indicates which variable in the returned |referenceVars| slice is used.
|
||||
func substituteVariableReferences(expression string, stack *InterpreterStack) (newExpression string, referencedVars []string, err error) {
|
||||
scanResult, err := pg_query.Scan(expression)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
|
||||
varMap := stack.ListVariables()
|
||||
for i := 0; i < len(scanResult.Tokens); i++ {
|
||||
token := scanResult.Tokens[i]
|
||||
substring := expression[token.Start:token.End]
|
||||
// varMap lowercases everything, so we'll lowercase our substring to enable case-insensitivity
|
||||
isAfterDot := i > 0 && scanResult.Tokens[i-1].Token == '.'
|
||||
|
||||
if !isAfterDot {
|
||||
if _, ok := varMap[strings.ToLower(substring)]; ok {
|
||||
// If there's a '.', then we'll assume this is accessing a record's field (`NEW.val1` for example)
|
||||
for i+2 < len(scanResult.Tokens) && scanResult.Tokens[i+1].Token == '.' {
|
||||
nextFieldSubstring := expression[scanResult.Tokens[i+2].Start:scanResult.Tokens[i+2].End]
|
||||
substring += "." + nextFieldSubstring
|
||||
i += 2
|
||||
}
|
||||
// Variables cannot have a '(' after their name as that would classify them as functions, so we have to
|
||||
// explicitly check for that. This is because variables and functions can share names, for example:
|
||||
// SELECT COUNT(*) INTO count FROM table_name;
|
||||
if i+1 >= len(scanResult.Tokens) || scanResult.Tokens[i+1].Token != '(' {
|
||||
referencedVars = append(referencedVars, substring)
|
||||
newExpression += fmt.Sprintf("$%d ", len(referencedVars))
|
||||
} else {
|
||||
newExpression += substring + " "
|
||||
}
|
||||
} else if _, ok := triggerSpecialVariables[substring]; ok {
|
||||
referencedVars = append(referencedVars, substring)
|
||||
newExpression += fmt.Sprintf("$%d ", len(referencedVars))
|
||||
} else {
|
||||
newExpression += substring + " "
|
||||
}
|
||||
} else {
|
||||
newExpression += substring + " "
|
||||
}
|
||||
}
|
||||
|
||||
return newExpression, referencedVars, nil
|
||||
}
|
||||
Reference in New Issue
Block a user