wehub-resource-sync
263 lines
9.2 KiB
Go
263 lines
9.2 KiB
Go
// Copyright 2024 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/go-mysql-server/sql/analyzer"
|
|
"github.com/dolthub/go-mysql-server/sql/expression/function"
|
|
"github.com/dolthub/vitess/go/vt/sqlparser"
|
|
|
|
"github.com/dolthub/doltgresql/postgres/parser/parser"
|
|
pgtypes "github.com/dolthub/doltgresql/server/types"
|
|
)
|
|
|
|
// Catalog contains all of the PostgreSQL functions.
|
|
var Catalog = map[string][]FunctionInterface{}
|
|
|
|
// AggregateCatalog contains all of the PostgreSQL aggregate functions.
|
|
var AggregateCatalog = map[string][]AggregateFunctionInterface{}
|
|
|
|
// initializedFunctions simply states whether Initialize has been called yet.
|
|
var initializedFunctions = false
|
|
|
|
// RegisterFunction registers the given function, so that it will be usable from a running server. This should be called
|
|
// from within an init().
|
|
func RegisterFunction(f FunctionInterface) {
|
|
if initializedFunctions {
|
|
panic("attempted to register a function after the init() phase")
|
|
}
|
|
switch f := f.(type) {
|
|
case Function0:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function1:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function1N:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function2:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function2N:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function3:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function4:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function5:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function6:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case Function7:
|
|
name := strings.ToLower(f.Name)
|
|
Catalog[name] = append(Catalog[name], f)
|
|
case InterpretedFunction:
|
|
name := strings.ToLower(f.ID.FunctionName())
|
|
Catalog[name] = append(Catalog[name], f)
|
|
default:
|
|
panic("unhandled function type")
|
|
}
|
|
}
|
|
|
|
// RegisterAggregateFunction registers the given function, so that it will be usable from a running server. This should be called
|
|
// from within an init().
|
|
func RegisterAggregateFunction(f AggregateFunctionInterface) {
|
|
if initializedFunctions {
|
|
panic("attempted to register a function after the init() phase")
|
|
}
|
|
switch f := f.(type) {
|
|
case Func1Aggregate:
|
|
name := strings.ToLower(f.Name)
|
|
AggregateCatalog[name] = append(AggregateCatalog[name], f)
|
|
default:
|
|
panic(fmt.Sprintf("unhandled function type %T", f))
|
|
}
|
|
}
|
|
|
|
// Initialize handles the initialization of the catalog by overwriting the built-in GMS functions, since they do not
|
|
// apply to PostgreSQL (and functions of the same name often have different behavior).
|
|
func Initialize(astConvert func(parser.Statement) (sqlparser.Statement, error)) {
|
|
// This should only be called once. We don't use sync.Once since we also want to panic if someone attempts to
|
|
// register a function after initialization.
|
|
if initializedFunctions {
|
|
return
|
|
}
|
|
initializedFunctions = true
|
|
|
|
convertToVitess = astConvert
|
|
pgtypes.LoadFunctionFromCatalog = getQuickFunctionForTypes
|
|
analyzer.ExternalFunctionProvider = &FunctionProvider{}
|
|
replaceGmsBuiltIns()
|
|
validateFunctions()
|
|
compileFunctions()
|
|
compileAggs()
|
|
}
|
|
|
|
// replaceGmsBuiltIns replaces all GMS built-ins that have conflicting names with PostgreSQL functions.
|
|
func replaceGmsBuiltIns() {
|
|
functionNames := make(map[string]struct{})
|
|
for name := range Catalog {
|
|
functionNames[strings.ToLower(name)] = struct{}{}
|
|
}
|
|
var newBuiltIns []sql.Function
|
|
for _, f := range function.BuiltIns {
|
|
if _, ok := functionNames[strings.ToLower(f.FunctionName())]; !ok {
|
|
newBuiltIns = append(newBuiltIns, f)
|
|
}
|
|
}
|
|
function.BuiltIns = newBuiltIns
|
|
}
|
|
|
|
// validateFunctions panics if any functions are defined incorrectly or ambiguously.
|
|
func validateFunctions() {
|
|
for funcName, overloads := range Catalog {
|
|
if err := validateFunction(funcName, overloads); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// validateFunction validates whether functions are defined incorrectly or ambiguously.
|
|
func validateFunction(funcName string, overloads []FunctionInterface) error {
|
|
// Verify that each function uses the correct Function overload
|
|
for _, functionOverload := range overloads {
|
|
if functionOverload.GetExpectedParameterCount() >= 0 &&
|
|
len(functionOverload.GetParameters()) != functionOverload.GetExpectedParameterCount() {
|
|
return errors.Errorf("function `%s` should have %d arguments but has %d arguments",
|
|
funcName, functionOverload.GetExpectedParameterCount(), len(functionOverload.GetParameters()))
|
|
}
|
|
}
|
|
// Verify that all overloads are unique
|
|
for functionIndex, f1 := range overloads {
|
|
for _, f2 := range overloads[functionIndex+1:] {
|
|
sameCount := 0
|
|
if f1.GetExpectedParameterCount() == f2.GetExpectedParameterCount() {
|
|
f2Parameters := f2.GetParameters()
|
|
for parameterIndex, f1Parameter := range f1.GetParameters() {
|
|
if f1Parameter.Equals(f2Parameters[parameterIndex]) {
|
|
sameCount++
|
|
}
|
|
}
|
|
}
|
|
if sameCount == f1.GetExpectedParameterCount() && f1.GetExpectedParameterCount() > 0 {
|
|
return errors.Errorf("duplicate function overloads on `%s`", funcName)
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// compileNonOperatorFunction creates a CompiledFunction for each overload of the given function.
|
|
func compileNonOperatorFunction(funcName string, overloads []FunctionInterface) {
|
|
overloadTree := NewOverloads()
|
|
for _, functionOverload := range overloads {
|
|
if err := overloadTree.Add(functionOverload); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
// Store the compiled function into the engine's built-in functions
|
|
// TODO: don't do this, use an actual contract for communicating these functions to the engine catalog
|
|
createFunc := func(ctx *sql.Context, params ...sql.Expression) (sql.Expression, error) {
|
|
return NewCompiledFunction(ctx, funcName, params, overloadTree, false), nil
|
|
}
|
|
function.BuiltIns = append(function.BuiltIns, sql.FunctionN{
|
|
Name: funcName,
|
|
Fn: createFunc,
|
|
})
|
|
compiledCatalog[funcName] = createFunc
|
|
}
|
|
|
|
// compileNonOperatorFunction creates a CompiledFunction for each overload of the given function.
|
|
func compileAggFunction(funcName string, overloads []AggregateFunctionInterface) {
|
|
var newBuffer NewBufferFn
|
|
overloadTree := NewOverloads()
|
|
for _, functionOverload := range overloads {
|
|
newBuffer = functionOverload.NewBuffer
|
|
if err := overloadTree.Add(functionOverload); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
// Store the compiled function into the engine's built-in functions
|
|
// TODO: don't do this, use an actual contract for communicating these functions to the engine catalog
|
|
createFunc := func(ctx *sql.Context, params ...sql.Expression) (sql.Expression, error) {
|
|
return NewCompiledAggregateFunction(ctx, funcName, params, overloadTree, newBuffer), nil
|
|
}
|
|
function.BuiltIns = append(function.BuiltIns, sql.FunctionN{
|
|
Name: funcName,
|
|
Fn: createFunc,
|
|
})
|
|
compiledCatalog[funcName] = createFunc
|
|
}
|
|
|
|
// compileFunctions creates a CompiledFunction for each overload of each function in the catalog.
|
|
func compileFunctions() {
|
|
for funcName, overloads := range Catalog {
|
|
compileNonOperatorFunction(funcName, overloads)
|
|
}
|
|
|
|
// Build the overload for all unary and binary functions based on their operator. This will be used for fallback if
|
|
// an exact match is not found. Compiled functions (which wrap the overload deducer) handle upcasting and other
|
|
// special rules, so it's far more efficient to reuse it for operators. Operators are also a special case since they
|
|
// all have different names, while standard overload deducers work on a function-name basis.
|
|
for signature, functionOverload := range unaryFunctions {
|
|
overloads, ok := unaryOperatorOverloads[signature.Operator]
|
|
if !ok {
|
|
overloads = NewOverloads()
|
|
unaryOperatorOverloads[signature.Operator] = overloads
|
|
}
|
|
if err := overloads.Add(functionOverload); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
for signature, functionOverload := range binaryFunctions {
|
|
overloads, ok := binaryOperatorOverloads[signature.Operator]
|
|
if !ok {
|
|
overloads = NewOverloads()
|
|
binaryOperatorOverloads[signature.Operator] = overloads
|
|
}
|
|
if err := overloads.Add(functionOverload); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
// Add all permutations for the unary and binary operators
|
|
for operator, overload := range unaryOperatorOverloads {
|
|
unaryOperatorPermutations[operator] = overload.overloadsForParams(1)
|
|
}
|
|
for operator, overload := range binaryOperatorOverloads {
|
|
binaryOperatorPermutations[operator] = overload.overloadsForParams(2)
|
|
}
|
|
}
|
|
|
|
func compileAggs() {
|
|
for funcName, overloads := range AggregateCatalog {
|
|
compileAggFunction(funcName, overloads)
|
|
}
|
|
}
|