wehub-resource-sync
333 lines
9.9 KiB
Go
333 lines
9.9 KiB
Go
// 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 framework
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/dolthub/go-mysql-server/sql"
|
|
"github.com/dolthub/vitess/go/vt/sqlparser"
|
|
|
|
"github.com/dolthub/doltgresql/core/id"
|
|
"github.com/dolthub/doltgresql/postgres/parser/parser"
|
|
"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
|
|
pgtypes "github.com/dolthub/doltgresql/server/types"
|
|
)
|
|
|
|
// SQLFunction is the implementation of functions created using SQL.
|
|
type SQLFunction struct {
|
|
ID id.Function
|
|
ReturnType *pgtypes.DoltgresType
|
|
ParameterNames []string
|
|
ParameterTypes []*pgtypes.DoltgresType
|
|
ParameterDefaults []string
|
|
Variadic bool
|
|
IsNonDeterministic bool
|
|
Strict bool
|
|
SqlStatement string
|
|
SetOf bool
|
|
ReturnTableType []*pgtypes.DoltgresType
|
|
}
|
|
|
|
var _ FunctionInterface = SQLFunction{}
|
|
|
|
// GetExpectedParameterCount implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) GetExpectedParameterCount() int {
|
|
return len(sqlFunc.ParameterTypes)
|
|
}
|
|
|
|
// GetName implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) GetName() string {
|
|
return sqlFunc.ID.FunctionName()
|
|
}
|
|
|
|
// GetParameters implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) GetParameters() []*pgtypes.DoltgresType {
|
|
return sqlFunc.ParameterTypes
|
|
}
|
|
|
|
// GetReturn implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) GetReturn() *pgtypes.DoltgresType {
|
|
return sqlFunc.ReturnType
|
|
}
|
|
|
|
// InternalID implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) InternalID() id.Id {
|
|
return sqlFunc.ID.AsId()
|
|
}
|
|
|
|
// IsStrict implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) IsStrict() bool {
|
|
return sqlFunc.Strict
|
|
}
|
|
|
|
// NonDeterministic implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) NonDeterministic() bool {
|
|
return sqlFunc.IsNonDeterministic
|
|
}
|
|
|
|
// IsCVariadic implements the FunctionInterface interface.
|
|
func (sqlFunc SQLFunction) IsCVariadic() bool {
|
|
return false
|
|
}
|
|
|
|
// VariadicIndex implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) VariadicIndex() int {
|
|
// TODO: implement variadic
|
|
return -1
|
|
}
|
|
|
|
// IsSRF implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) IsSRF() bool {
|
|
return sqlFunc.SetOf
|
|
}
|
|
|
|
// enforceInterfaceInheritance implements the interface FunctionInterface.
|
|
func (sqlFunc SQLFunction) enforceInterfaceInheritance(error) {}
|
|
|
|
// CallSqlFunction runs the given SQL definition inside the function on the given runner.
|
|
func CallSqlFunction(ctx *sql.Context, f SQLFunction, runner sql.StatementRunner, args []any) (any, error) {
|
|
paramMap := make(map[string]*ParamTypAndValue)
|
|
for i, name := range f.ParameterNames {
|
|
if name == "" {
|
|
// This allows for positional references such as $1, $2, etc.
|
|
name = fmt.Sprintf("$%d", i+1)
|
|
}
|
|
paramMap[name] = &ParamTypAndValue{
|
|
Typ: f.ParameterTypes[i],
|
|
Val: args[i],
|
|
FromCreate: false,
|
|
}
|
|
}
|
|
|
|
if lower := strings.ToLower(f.SqlStatement); strings.HasPrefix(lower, "return") {
|
|
f.SqlStatement = fmt.Sprintf("SELECT%s", f.SqlStatement[6:])
|
|
}
|
|
|
|
parseds, err := parser.Parse(f.SqlStatement)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
if len(parseds) > 1 {
|
|
// of multiple statements, the function returns the result of the final statement in the execution block
|
|
var res any
|
|
for _, parsed := range parseds {
|
|
err = ReplaceFunctionColumn(parsed.AST, paramMap)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
convertedAST, err := convertToVitess(parsed)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
res, err = sql.RunInterpreted(ctx, func(subCtx *sql.Context) (any, error) {
|
|
sch, rowIter, _, err := runner.QueryWithBindings(ctx, parsed.AST.String(), convertedAST, nil, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
rows, err := sql.RowIterToRows(subCtx, rowIter)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(sch) != 1 {
|
|
return nil, errors.New("expression does not result in a single value")
|
|
}
|
|
if len(rows) != 1 {
|
|
return nil, errors.New("expression returned multiple result sets")
|
|
}
|
|
if len(rows[0]) != 1 {
|
|
return nil, errors.New("expression returned multiple results")
|
|
}
|
|
return rows[0][0], nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
if f.ReturnType.ID == pgtypes.Void.ID {
|
|
return nil, nil
|
|
}
|
|
return res, nil
|
|
}
|
|
|
|
// single statement
|
|
parsed := parseds[0]
|
|
err = ReplaceFunctionColumn(parsed.AST, paramMap)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
convertedAST, err := convertToVitess(parsed)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// stmt.AST is updated at this point with FunctionColumn
|
|
return sql.RunInterpreted(ctx, func(subCtx *sql.Context) (any, error) {
|
|
sch, rowIter, _, err := runner.QueryWithBindings(ctx, parsed.AST.String(), convertedAST, nil, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if !f.SetOf {
|
|
rows, err := sql.RowIterToRows(subCtx, rowIter)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// single column row result
|
|
if len(sch) != 1 {
|
|
return nil, errors.New("expression does not result in a single value")
|
|
}
|
|
if len(rows) != 1 {
|
|
return nil, errors.New("expression returned multiple result sets")
|
|
}
|
|
if len(rows[0]) != 1 {
|
|
return nil, errors.New("expression returned multiple results")
|
|
}
|
|
return rows[0][0], nil
|
|
}
|
|
// multiple column row result
|
|
if f.ReturnType.TypCategory == pgtypes.TypeCategory_CompositeTypes {
|
|
// record type
|
|
return rowIterToRecord(ctx, rowIter, sch)
|
|
}
|
|
return rowIter, nil
|
|
})
|
|
}
|
|
|
|
// ParamTypAndValue contains the parameter type and
|
|
// string value of argument if applicable
|
|
type ParamTypAndValue struct {
|
|
Typ *pgtypes.DoltgresType
|
|
Val any
|
|
FromCreate bool
|
|
}
|
|
|
|
// ReplaceFunctionColumn parses and replaces UnresolvedName and Placeholder expressions
|
|
// with FunctionColumn expression containing parameter type and arguments if applicable.
|
|
// It also replaces empty parameter name with binding variable name to match the name used in FunctionColumn.
|
|
// This function should be used for FUNCTION with SQL language statements only.
|
|
func ReplaceFunctionColumn(parsedAST tree.Statement, params map[string]*ParamTypAndValue) error {
|
|
if len(params) == 0 {
|
|
return nil
|
|
}
|
|
// Function's final statement must be SELECT or INSERT/UPDATE/DELETE RETURNING
|
|
switch s := parsedAST.(type) {
|
|
case *tree.Select:
|
|
switch t := s.Select.(type) {
|
|
case *tree.SelectClause:
|
|
for i, e := range t.Exprs {
|
|
t.Exprs[i].Expr = ReplaceUnresolvedToFunctionColumn(params, e.Expr)
|
|
}
|
|
if t.Where != nil {
|
|
t.Where.Expr = ReplaceUnresolvedToFunctionColumn(params, t.Where.Expr)
|
|
}
|
|
case *tree.ValuesClause:
|
|
for i, row := range t.Rows {
|
|
for j, e := range row {
|
|
row[j] = ReplaceUnresolvedToFunctionColumn(params, e)
|
|
}
|
|
t.Rows[i] = row
|
|
}
|
|
}
|
|
return nil
|
|
case *tree.Insert:
|
|
err := ReplaceFunctionColumn(s.Rows, params)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
case *tree.Update:
|
|
for i, e := range s.Exprs {
|
|
s.Exprs[i].Expr = ReplaceUnresolvedToFunctionColumn(params, e.Expr)
|
|
}
|
|
if s.Where != nil {
|
|
s.Where.Expr = ReplaceUnresolvedToFunctionColumn(params, s.Where.Expr)
|
|
}
|
|
return nil
|
|
case *tree.Delete:
|
|
if s.Where != nil {
|
|
s.Where.Expr = ReplaceUnresolvedToFunctionColumn(params, s.Where.Expr)
|
|
}
|
|
return nil
|
|
case *tree.Truncate:
|
|
return nil
|
|
case *tree.Return:
|
|
s.Expr = ReplaceUnresolvedToFunctionColumn(params, s.Expr)
|
|
return nil
|
|
default:
|
|
return errors.Errorf("unsupported statement defined in function: %T", parsedAST)
|
|
}
|
|
}
|
|
|
|
// ReplaceUnresolvedToFunctionColumn replaces Placeholder and UnresolvedName expressions with FunctionColumn containing
|
|
// parameter type and argument value if applicable when the name of expression matches function parameter.
|
|
func ReplaceUnresolvedToFunctionColumn(paramMap map[string]*ParamTypAndValue, expr tree.Expr) tree.Expr {
|
|
e, _ := tree.SimpleVisit(expr, func(visitingExpr tree.Expr) (recurse bool, newExpr tree.Expr, err error) {
|
|
switch v := visitingExpr.(type) {
|
|
case *tree.Placeholder:
|
|
name := fmt.Sprintf("$%d", v.Idx+1)
|
|
if tv, ok := paramMap[name]; ok {
|
|
return false, tree.FunctionColumn{
|
|
Name: name,
|
|
Typ: tv.Typ,
|
|
Idx: uint16(v.Idx),
|
|
FromCreate: tv.FromCreate,
|
|
Val: tv.Val,
|
|
}, nil
|
|
}
|
|
case *tree.UnresolvedName:
|
|
name := v.String()
|
|
if tv, ok := paramMap[name]; ok {
|
|
return false, tree.FunctionColumn{
|
|
Name: name,
|
|
Typ: tv.Typ,
|
|
FromCreate: tv.FromCreate,
|
|
Val: tv.Val,
|
|
}, nil
|
|
}
|
|
}
|
|
return true, visitingExpr, nil
|
|
})
|
|
return e
|
|
}
|
|
|
|
// rowIterToRecord converts given rows with schema provided to rowIter containing array of pgtypes.RecordValue.
|
|
func rowIterToRecord(ctx *sql.Context, rowIter sql.RowIter, sch sql.Schema) (sql.RowIter, error) {
|
|
rows, err := sql.RowIterToRows(ctx, rowIter)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var newRows = make([]sql.Row, len(rows))
|
|
for i, row := range rows {
|
|
if len(row) != len(sch) {
|
|
return nil, errors.New("number of row values does not match number of schema columns")
|
|
}
|
|
var r = make([]pgtypes.RecordValue, len(sch))
|
|
for j, col := range sch {
|
|
r[j] = pgtypes.RecordValue{
|
|
Type: col.Type.(*pgtypes.DoltgresType),
|
|
Value: row[j],
|
|
}
|
|
}
|
|
newRows[i] = sql.Row{r}
|
|
}
|
|
return sql.RowsToRowIter(newRows...), nil
|
|
}
|
|
|
|
// convertToVitess is set by init and is used to avoid import cycles
|
|
var convertToVitess func(postgresStmt parser.Statement) (sqlparser.Statement, error)
|