chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,343 @@
|
||||
// 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 functions
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
|
||||
"github.com/dolthub/doltgresql/core"
|
||||
"github.com/dolthub/doltgresql/core/id"
|
||||
"github.com/dolthub/doltgresql/server/functions/framework"
|
||||
pgtypes "github.com/dolthub/doltgresql/server/types"
|
||||
"github.com/dolthub/doltgresql/utils"
|
||||
)
|
||||
|
||||
// initArray registers the functions to the catalog.
|
||||
func initArray() {
|
||||
framework.RegisterFunction(array_in)
|
||||
framework.RegisterFunction(array_out)
|
||||
framework.RegisterFunction(array_recv)
|
||||
framework.RegisterFunction(array_send)
|
||||
framework.RegisterFunction(btarraycmp)
|
||||
framework.RegisterFunction(array_subscript_handler)
|
||||
}
|
||||
|
||||
// array_in represents the PostgreSQL function of array type IO input.
|
||||
var array_in = framework.Function3{
|
||||
Name: "array_in",
|
||||
Return: pgtypes.AnyArray,
|
||||
Parameters: [3]*pgtypes.DoltgresType{pgtypes.Cstring, pgtypes.Oid, pgtypes.Int32},
|
||||
Strict: true,
|
||||
Callable: func(ctx *sql.Context, _ [4]*pgtypes.DoltgresType, val1, val2, val3 any) (any, error) {
|
||||
input := val1.(string)
|
||||
baseTypeOid := val2.(id.Id)
|
||||
baseType := pgtypes.IDToBuiltInDoltgresType[id.Type(baseTypeOid)]
|
||||
if baseType == nil {
|
||||
if typColl, err := pgtypes.GetTypesCollectionFromContext(ctx, ""); err == nil && typColl != nil {
|
||||
if t, err := typColl.GetType(ctx, id.Type(baseTypeOid)); err == nil {
|
||||
baseType = t
|
||||
}
|
||||
}
|
||||
}
|
||||
if baseType == nil {
|
||||
return nil, errors.Errorf("unknown array element type: %s", string(baseTypeOid))
|
||||
}
|
||||
typmod := val3.(int32)
|
||||
baseType = baseType.WithAttTypMod(typmod)
|
||||
if len(input) < 2 || input[0] != '{' || input[len(input)-1] != '}' {
|
||||
// This error is regarded as a critical error, and thus we immediately return the error alongside a nil
|
||||
// value. Returning a nil value is a signal to not ignore the error.
|
||||
return nil, errors.Errorf(`malformed array literal: "%s"`, input)
|
||||
}
|
||||
// We'll remove the surrounding braces since we've already verified that they're there
|
||||
input = input[1 : len(input)-1]
|
||||
var values []any
|
||||
var err error
|
||||
sb := strings.Builder{}
|
||||
quoteStartCount := 0
|
||||
quoteEndCount := 0
|
||||
escaped := false
|
||||
// Iterate over each rune in the input to collect and process the rune elements
|
||||
for _, r := range input {
|
||||
if escaped {
|
||||
sb.WriteRune(r)
|
||||
escaped = false
|
||||
} else if quoteStartCount > quoteEndCount {
|
||||
switch r {
|
||||
case '\\':
|
||||
escaped = true
|
||||
case '"':
|
||||
quoteEndCount++
|
||||
default:
|
||||
sb.WriteRune(r)
|
||||
}
|
||||
} else {
|
||||
switch r {
|
||||
case ' ', '\t', '\n', '\r':
|
||||
continue
|
||||
case '\\':
|
||||
escaped = true
|
||||
case '"':
|
||||
quoteStartCount++
|
||||
case ',':
|
||||
if quoteStartCount >= 2 {
|
||||
// This is a malformed string, thus we treat it as a critical error.
|
||||
return nil, errors.Errorf(`malformed array literal: "%s"`, input)
|
||||
}
|
||||
str := sb.String()
|
||||
var innerValue any
|
||||
if quoteStartCount == 0 && strings.EqualFold(str, "null") {
|
||||
// An unquoted case-insensitive NULL is treated as an actual null value
|
||||
innerValue = nil
|
||||
} else {
|
||||
var nErr error
|
||||
innerValue, nErr = baseType.IoInput(ctx, str)
|
||||
if nErr != nil && err == nil {
|
||||
// This is a non-critical error, therefore the error may be ignored at a higher layer (such as
|
||||
// an explicit cast) and the inner type will still return a valid result, so we must allow the
|
||||
// values to propagate.
|
||||
err = nErr
|
||||
}
|
||||
}
|
||||
values = append(values, innerValue)
|
||||
sb.Reset()
|
||||
quoteStartCount = 0
|
||||
quoteEndCount = 0
|
||||
default:
|
||||
sb.WriteRune(r)
|
||||
}
|
||||
}
|
||||
}
|
||||
// Use anything remaining in the buffer as the last element
|
||||
if sb.Len() > 0 {
|
||||
if escaped || quoteStartCount > quoteEndCount || quoteStartCount >= 2 {
|
||||
// These errors are regarded as critical errors, and thus we immediately return the error alongside a nil
|
||||
// value. Returning a nil value is a signal to not ignore the error.
|
||||
return nil, errors.Errorf(`malformed array literal: "%s"`, input)
|
||||
} else {
|
||||
str := sb.String()
|
||||
var innerValue any
|
||||
if quoteStartCount == 0 && strings.EqualFold(str, "NULL") {
|
||||
// An unquoted case-insensitive NULL is treated as an actual null value
|
||||
innerValue = nil
|
||||
} else {
|
||||
var nErr error
|
||||
innerValue, nErr = baseType.IoInput(ctx, str)
|
||||
if nErr != nil && err == nil {
|
||||
// This is a non-critical error, therefore the error may be ignored at a higher layer (such as
|
||||
// an explicit cast) and the inner type will still return a valid result, so we must allow the
|
||||
// values to propagate.
|
||||
err = nErr
|
||||
}
|
||||
}
|
||||
values = append(values, innerValue)
|
||||
}
|
||||
}
|
||||
|
||||
return values, err
|
||||
},
|
||||
}
|
||||
|
||||
// array_out represents the PostgreSQL function of array type IO output.
|
||||
var array_out = framework.Function1{
|
||||
Name: "array_out",
|
||||
Return: pgtypes.Cstring,
|
||||
Parameters: [1]*pgtypes.DoltgresType{pgtypes.AnyArray},
|
||||
Strict: true,
|
||||
Callable: func(ctx *sql.Context, t [2]*pgtypes.DoltgresType, val any) (any, error) {
|
||||
arrType := t[0]
|
||||
baseType := arrType.ArrayBaseType()
|
||||
return pgtypes.ArrToString(ctx, val.([]any), baseType, false)
|
||||
},
|
||||
}
|
||||
|
||||
// array_recv represents the PostgreSQL function of array type IO receive.
|
||||
var array_recv = framework.Function3{
|
||||
Name: "array_recv",
|
||||
Return: pgtypes.AnyArray,
|
||||
Parameters: [3]*pgtypes.DoltgresType{pgtypes.Internal, pgtypes.Oid, pgtypes.Int32},
|
||||
Strict: true,
|
||||
Callable: array_recv_callable,
|
||||
}
|
||||
|
||||
// array_recv_callable is the function definition of array_recv.
|
||||
func array_recv_callable(ctx *sql.Context, t [4]*pgtypes.DoltgresType, val1, val2, val3 any) (any, error) {
|
||||
data := val1.([]byte)
|
||||
if data == nil {
|
||||
return nil, nil
|
||||
}
|
||||
typeColl, err := core.GetTypesCollectionFromContext(ctx, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
reader := utils.NewWireReader(data)
|
||||
dimensions := reader.ReadInt32()
|
||||
_ = reader.ReadInt32() // Whether the array has a null, doesn't seem useful
|
||||
baseTypeID := id.Type(id.Cache().ToInternal(reader.ReadUint32()))
|
||||
baseType, err := typeColl.GetType(ctx, baseTypeID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if baseType == nil {
|
||||
return nil, pgtypes.ErrTypeDoesNotExist.New(baseTypeID.TypeName())
|
||||
}
|
||||
// TODO: handle more than 1 dimension
|
||||
if dimensions > 1 {
|
||||
return nil, errors.Errorf("array dimensions greater than 1 are not yet supported")
|
||||
}
|
||||
var vals []any
|
||||
for dimensionIdx := int32(0); dimensionIdx < dimensions; dimensionIdx++ {
|
||||
elementsCount := reader.ReadInt32()
|
||||
_ = reader.ReadInt32() // Lower bound, not sure what to do with this
|
||||
for i := int32(0); i < elementsCount; i++ {
|
||||
elementLen := reader.ReadInt32()
|
||||
if elementLen != -1 {
|
||||
valBytes := reader.ReadBytes(uint32(elementLen))
|
||||
val, err := baseType.CallReceive(ctx, valBytes)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
vals = append(vals, val)
|
||||
} else {
|
||||
vals = append(vals, nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
return vals, nil
|
||||
}
|
||||
|
||||
// array_send represents the PostgreSQL function of array type IO send.
|
||||
var array_send = framework.Function1{
|
||||
Name: "array_send",
|
||||
Return: pgtypes.Bytea,
|
||||
Parameters: [1]*pgtypes.DoltgresType{pgtypes.AnyArray},
|
||||
Strict: true,
|
||||
Callable: func(ctx *sql.Context, t [2]*pgtypes.DoltgresType, val any) (any, error) {
|
||||
if wrapper, ok := val.(sql.AnyWrapper); ok {
|
||||
var err error
|
||||
val, err = wrapper.UnwrapAny(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if val == nil {
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
vals := val.([]any)
|
||||
// Check for nulls first
|
||||
hasNull := false
|
||||
for _, val := range vals {
|
||||
if val == nil {
|
||||
hasNull = true
|
||||
break
|
||||
}
|
||||
}
|
||||
// Count the number of dimensions
|
||||
dimensions := int32(0)
|
||||
innerVals := vals
|
||||
for len(innerVals) > 0 {
|
||||
dimensions++
|
||||
slice, ok := innerVals[0].([]any)
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
innerVals = slice
|
||||
}
|
||||
if dimensions > 1 {
|
||||
return nil, errors.Errorf("arrays with %d dimensions are not yet supported using the binary format", dimensions)
|
||||
}
|
||||
writer := utils.NewWireWriter()
|
||||
writer.WriteInt32(dimensions) // Write the number of dimensions
|
||||
if hasNull {
|
||||
writer.WriteInt32(1)
|
||||
} else {
|
||||
writer.WriteInt32(0)
|
||||
}
|
||||
writer.WriteUint32(id.Cache().ToOID(t[0].BaseType().ID.AsId())) // Element OID
|
||||
for i := int32(0); i < dimensions; i++ {
|
||||
writer.WriteInt32(int32(len(vals))) // Elements in this dimension
|
||||
if t[0].IsArrayType() {
|
||||
writer.WriteInt32(1) // Lower bound, or what index number we start at (seems to always be 1?)
|
||||
} else {
|
||||
writer.WriteInt32(0)
|
||||
}
|
||||
for _, val := range vals {
|
||||
if val == nil {
|
||||
writer.WriteInt32(-1)
|
||||
} else {
|
||||
valBytes, err := t[0].BaseType().CallSend(ctx, val)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
writer.WriteInt32(int32(len(valBytes)))
|
||||
writer.WriteBytes(valBytes)
|
||||
}
|
||||
}
|
||||
}
|
||||
return writer.BufferData(), nil
|
||||
},
|
||||
}
|
||||
|
||||
// btarraycmp represents the PostgreSQL function of array type byte compare.
|
||||
var btarraycmp = framework.Function2{
|
||||
Name: "btarraycmp",
|
||||
Return: pgtypes.Int32,
|
||||
Parameters: [2]*pgtypes.DoltgresType{pgtypes.AnyArray, pgtypes.AnyArray},
|
||||
Strict: true,
|
||||
Callable: func(ctx *sql.Context, t [3]*pgtypes.DoltgresType, val1, val2 any) (any, error) {
|
||||
at := t[0]
|
||||
bt := t[1]
|
||||
if !at.Equals(bt) {
|
||||
// TODO: currently, types should match.
|
||||
// Technically, does not have to e.g.: float4 vs float8
|
||||
return nil, errors.Errorf("different type comparison is not supported yet")
|
||||
}
|
||||
|
||||
ab := val1.([]any)
|
||||
bb := val2.([]any)
|
||||
minLength := utils.Min(len(ab), len(bb))
|
||||
for i := 0; i < minLength; i++ {
|
||||
res, err := at.ArrayBaseType().Compare(ctx, ab[i], bb[i])
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if res != 0 {
|
||||
return res, nil
|
||||
}
|
||||
}
|
||||
if len(ab) == len(bb) {
|
||||
return int32(0), nil
|
||||
} else if len(ab) < len(bb) {
|
||||
return int32(-1), nil
|
||||
} else {
|
||||
return int32(1), nil
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
// array_subscript_handler represents the PostgreSQL function of array type subscript handler.
|
||||
var array_subscript_handler = framework.Function1{
|
||||
Name: "array_subscript_handler",
|
||||
Return: pgtypes.Internal,
|
||||
Parameters: [1]*pgtypes.DoltgresType{pgtypes.Internal},
|
||||
Strict: true,
|
||||
Callable: func(ctx *sql.Context, t [2]*pgtypes.DoltgresType, val any) (any, error) {
|
||||
// TODO
|
||||
return []byte{}, nil
|
||||
},
|
||||
}
|
||||
Reference in New Issue
Block a user