Files
2026-07-13 12:32:25 +08:00

290 lines
9.9 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 types
import (
"cmp"
"maps"
"slices"
"github.com/cockroachdb/errors"
"github.com/dolthub/go-mysql-server/sql"
"github.com/dolthub/go-mysql-server/sql/types"
"github.com/dolthub/doltgresql/core/id"
"github.com/dolthub/doltgresql/utils"
)
// init sets the serialization and deserialization functions.
func init() {
types.SetExtendedTypeSerializers(SerializeType, DeserializeType)
}
// SerializeType is able to serialize the given extended type into a byte slice. All extended types will be defined
// by DoltgreSQL.
func SerializeType(ctx *sql.Context, extendedType sql.ExtendedType) ([]byte, error) {
if doltgresType, ok := extendedType.(*DoltgresType); ok {
if doltgresType.IsUnresolved {
return nil, errors.Errorf(`attempted to serialize the unresolved type: %s`, doltgresType.Name())
}
return doltgresType.Serialize(), nil
}
return nil, errors.Errorf("unknown type to serialize")
}
// DeserializeType is able to deserialize the given serialized type into an appropriate extended type. All extended
// types will be defined by DoltgreSQL.
func DeserializeType(ctx *sql.Context, serializedType []byte) (sql.ExtendedType, error) {
if len(serializedType) == 0 {
return nil, errors.Errorf("deserializing empty type data")
}
typ := &DoltgresType{}
var typeColl TypeCollection
reader := utils.NewReader(serializedType)
version := reader.VariableUint()
if version != 0 {
return nil, errors.Errorf("version %d of types is not supported, please upgrade the server", version)
}
typ.ID = id.Type(reader.Id())
typ.TypLength = reader.Int16()
typ.PassedByVal = reader.Bool()
typ.TypType = TypeType(reader.String())
typ.TypCategory = TypeCategory(reader.String())
typ.IsPreferred = reader.Bool()
typ.IsDefined = reader.Bool()
typ.Delimiter = reader.String()
typ.RelID = reader.Id()
typ.SubscriptFunc = globalFunctionRegistry.InternalToRegistryID(id.Function(reader.Id()))
typ.Elem = prefillTypeDuringDeserialization(id.Type(reader.Id()))
typ.Array = prefillTypeDuringDeserialization(id.Type(reader.Id()))
typ.InputFunc = globalFunctionRegistry.InternalToRegistryID(id.Function(reader.Id()))
typ.OutputFunc = globalFunctionRegistry.InternalToRegistryID(id.Function(reader.Id()))
receiveFunc := id.Function(reader.Id())
typ.ReceiveFunc = globalFunctionRegistry.InternalToRegistryID(receiveFunc)
sendFunc := id.Function(reader.Id())
typ.SendFunc = globalFunctionRegistry.InternalToRegistryID(sendFunc)
typ.ModInFunc = globalFunctionRegistry.InternalToRegistryID(id.Function(reader.Id()))
typ.ModOutFunc = globalFunctionRegistry.InternalToRegistryID(id.Function(reader.Id()))
typ.AnalyzeFunc = globalFunctionRegistry.InternalToRegistryID(id.Function(reader.Id()))
typ.Align = TypeAlignment(reader.String())
typ.Storage = TypeStorage(reader.String())
typ.NotNull = reader.Bool()
typ.BaseTypeType = prefillTypeDuringDeserialization(id.Type(reader.Id()))
typ.TypMod = reader.Int32()
typ.NDims = reader.Int32()
typ.TypCollation = id.Collation(reader.Id())
typ.DefaulBin = reader.String()
typ.Default = reader.String()
numOfAcl := reader.VariableUint()
for k := uint64(0); k < numOfAcl; k++ {
ac := reader.String()
typ.Acl = append(typ.Acl, ac)
}
numOfChecks := reader.VariableUint()
for k := uint64(0); k < numOfChecks; k++ {
checkName := reader.String()
checkExpr := reader.String()
typ.Checks = append(typ.Checks, &sql.CheckDefinition{
Name: checkName,
CheckExpression: checkExpr,
Enforced: true,
})
}
typ.attTypMod = reader.Int32()
typ.CompareFunc = globalFunctionRegistry.InternalToRegistryID(id.Function(reader.Id()))
numOfEnumLabels := reader.VariableUint()
if numOfEnumLabels > 0 {
typ.EnumLabels = make(map[string]EnumLabel)
for k := uint64(0); k < numOfEnumLabels; k++ {
typeID := reader.Id()
sortOrder := reader.Float32()
typ.EnumLabels[typeID.Segment(1)] = EnumLabel{
ID: id.EnumLabel(typeID),
SortOrder: sortOrder,
}
}
}
numOfCompAttrs := reader.VariableUint()
if numOfCompAttrs > 0 {
typ.CompositeAttrs = make([]CompositeAttribute, numOfCompAttrs)
for k := uint64(0); k < numOfCompAttrs; k++ {
relID := reader.Id()
name := reader.String()
attrType := prefillTypeDuringDeserialization(id.Type(reader.Id()))
num := reader.Int16()
collation := reader.String()
typ.CompositeAttrs[k] = CompositeAttribute{
RelID: relID,
Name: name,
Type: attrType,
Num: num,
Collation: collation,
}
}
}
typ.InternalName = reader.String()
if !reader.IsEmpty() {
return nil, errors.Errorf("extra data found while deserializing type %s", typ.Name())
}
// Grab the serialization and deserialization functions for the built-in types
if f, ok := idToInternalSerializationFunc[sendFunc]; ok {
typ.SerializationFunc = f
}
if f, ok := idToInternalDeserializationFunc[receiveFunc]; ok {
typ.DeserializationFunc = f
}
// Resolve all of the potentially unresolved types
var err error
typ.Elem, typeColl, err = recursiveDeserializeType(ctx, typ, typeColl, typ.Elem)
if err != nil {
return nil, err
}
typ.Array, typeColl, err = recursiveDeserializeType(ctx, typ, typeColl, typ.Array)
if err != nil {
return nil, err
}
typ.BaseTypeType, typeColl, err = recursiveDeserializeType(ctx, typ, typeColl, typ.BaseTypeType)
if err != nil {
return nil, err
}
for i := range typ.CompositeAttrs {
typ.CompositeAttrs[i].Type, typeColl, err = recursiveDeserializeType(ctx, typ, typeColl, typ.CompositeAttrs[i].Type)
if err != nil {
return nil, err
}
}
// Return the deserialized object
return typ, nil
}
// Serialize returns the DoltgresType as a byte slice.
func (t *DoltgresType) Serialize() []byte {
writer := utils.NewWriter(256)
writer.VariableUint(0) // Version
// Write the type to the writer
writer.Id(t.ID.AsId())
writer.Int16(t.TypLength)
writer.Bool(t.PassedByVal)
writer.String(string(t.TypType))
writer.String(string(t.TypCategory))
writer.Bool(t.IsPreferred)
writer.Bool(t.IsDefined)
writer.String(t.Delimiter)
writer.Id(t.RelID)
writer.Id(globalFunctionRegistry.GetInternalID(t.SubscriptFunc).AsId())
writer.Id(t.Elem.ID.AsId())
writer.Id(t.Array.ID.AsId())
writer.Id(globalFunctionRegistry.GetInternalID(t.InputFunc).AsId())
writer.Id(globalFunctionRegistry.GetInternalID(t.OutputFunc).AsId())
writer.Id(globalFunctionRegistry.GetInternalID(t.ReceiveFunc).AsId())
writer.Id(globalFunctionRegistry.GetInternalID(t.SendFunc).AsId())
writer.Id(globalFunctionRegistry.GetInternalID(t.ModInFunc).AsId())
writer.Id(globalFunctionRegistry.GetInternalID(t.ModOutFunc).AsId())
writer.Id(globalFunctionRegistry.GetInternalID(t.AnalyzeFunc).AsId())
writer.String(string(t.Align))
writer.String(string(t.Storage))
writer.Bool(t.NotNull)
writer.Id(t.BaseTypeType.ID.AsId())
writer.Int32(t.TypMod)
writer.Int32(t.NDims)
writer.Id(t.TypCollation.AsId())
writer.String(t.DefaulBin)
writer.String(t.Default)
writer.VariableUint(uint64(len(t.Acl)))
for _, ac := range t.Acl {
writer.String(ac)
}
writer.VariableUint(uint64(len(t.Checks)))
for _, check := range t.Checks {
writer.String(check.Name)
writer.String(check.CheckExpression)
}
writer.Int32(t.attTypMod)
writer.Id(globalFunctionRegistry.GetInternalID(t.CompareFunc).AsId())
writer.VariableUint(uint64(len(t.EnumLabels)))
if len(t.EnumLabels) > 0 {
labels := slices.SortedFunc(maps.Values(t.EnumLabels), func(v1 EnumLabel, v2 EnumLabel) int {
return cmp.Compare(v1.ID, v2.ID)
})
for _, l := range labels {
writer.Id(l.ID.AsId())
writer.Float32(l.SortOrder)
}
}
writer.VariableUint(uint64(len(t.CompositeAttrs)))
if len(t.CompositeAttrs) > 0 {
for _, l := range t.CompositeAttrs {
writer.Id(l.RelID)
writer.String(l.Name)
writer.Id(l.Type.ID.AsId())
writer.Int16(l.Num)
writer.String(l.Collation)
}
}
writer.String(t.InternalName)
return writer.Data()
}
// prefillTypeDuringDeserialization attempts to return a suitable type if it's one that relies strictly on built-in
// types.
func prefillTypeDuringDeserialization(target id.Type) *DoltgresType {
if target == id.NullType {
return internalNullType
}
if builtin, ok := IDToBuiltInDoltgresType[target]; ok {
return builtin
}
return NewUnresolvedDoltgresTypeFromID(target)
}
// recursiveDeserializeType handles recursive type initialization by caching the type-in-progress and returning the
// target type. This also takes a TypeCollection so that it may be reused if multiple calls are made without having to
// fetch it at each instance.
func recursiveDeserializeType(ctx *sql.Context, typ *DoltgresType, typeColl TypeCollection, target *DoltgresType) (*DoltgresType, TypeCollection, error) {
if !target.IsUnresolved {
return target, typeColl, nil
}
// If there is no Doltgres context/session yet, then we can't load types from the context.
// This can happen when Dolt's AutoIncrementTracker scans tables at db-load time.
if ctx == nil {
return target, typeColl, nil
}
var recursedType *DoltgresType
var err error
if typeColl == nil {
typeColl, err = GetTypesCollectionFromContext(ctx, "")
if err != nil {
return nil, nil, err
}
}
typeColl.WithCachedType(typ, func() {
t, nErr := typeColl.GetType(ctx, target.ID)
recursedType = t
err = nErr
})
// `GetType` returns a nil type if it cannot be found, so we want to error in that case
if err == nil && recursedType == nil {
return target, typeColl, errors.Errorf("unable to resolve type `%s` during deserialization", target.ID.TypeName())
}
return recursedType, typeColl, err
}