chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,236 @@
|
||||
// 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 server
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
"github.com/dolthub/go-mysql-server/sql/expression"
|
||||
"github.com/dolthub/go-mysql-server/sql/plan"
|
||||
"github.com/dolthub/go-mysql-server/sql/transform"
|
||||
vitess "github.com/dolthub/vitess/go/vt/sqlparser"
|
||||
"github.com/jackc/pgx/v5/pgproto3"
|
||||
"github.com/lib/pq/oid"
|
||||
|
||||
"github.com/dolthub/doltgresql/core/dataloader"
|
||||
"github.com/dolthub/doltgresql/core/id"
|
||||
pgexprs "github.com/dolthub/doltgresql/server/expression"
|
||||
"github.com/dolthub/doltgresql/server/functions/framework"
|
||||
"github.com/dolthub/doltgresql/server/node"
|
||||
pgtypes "github.com/dolthub/doltgresql/server/types"
|
||||
)
|
||||
|
||||
// ErrorResponseSeverity represents the severity of an ErrorResponse message.
|
||||
type ErrorResponseSeverity string
|
||||
|
||||
const (
|
||||
ErrorResponseSeverity_Error ErrorResponseSeverity = "ERROR"
|
||||
ErrorResponseSeverity_Fatal ErrorResponseSeverity = "FATAL"
|
||||
ErrorResponseSeverity_Panic ErrorResponseSeverity = "PANIC"
|
||||
ErrorResponseSeverity_Warning ErrorResponseSeverity = "WARNING"
|
||||
ErrorResponseSeverity_Notice ErrorResponseSeverity = "NOTICE"
|
||||
ErrorResponseSeverity_Debug ErrorResponseSeverity = "DEBUG"
|
||||
ErrorResponseSeverity_Info ErrorResponseSeverity = "INFO"
|
||||
ErrorResponseSeverity_Log ErrorResponseSeverity = "LOG"
|
||||
)
|
||||
|
||||
// ReadyForQueryTransactionIndicator indicates the state of the transaction related to the query.
|
||||
type ReadyForQueryTransactionIndicator byte
|
||||
|
||||
const (
|
||||
ReadyForQueryTransactionIndicator_Idle ReadyForQueryTransactionIndicator = 'I'
|
||||
ReadyForQueryTransactionIndicator_TransactionBlock ReadyForQueryTransactionIndicator = 'T'
|
||||
ReadyForQueryTransactionIndicator_FailedTransactionBlock ReadyForQueryTransactionIndicator = 'E'
|
||||
)
|
||||
|
||||
// ConvertedQuery represents a query that has been converted from the Postgres representation to the Vitess
|
||||
// representation. String may contain the string version of the converted query. AST will contain the tree
|
||||
// version of the converted query, and is the recommended form to use. If AST is nil, then use the String version,
|
||||
// otherwise always prefer to AST.
|
||||
type ConvertedQuery struct {
|
||||
String string
|
||||
AST vitess.Statement
|
||||
StatementTag string
|
||||
}
|
||||
|
||||
// copyFromStdinState tracks the metadata for an import of data into a table using a COPY FROM STDIN statement. When
|
||||
// this statement is processed, the server accepts COPY DATA messages from the client with chunks of data to load
|
||||
// into a table.
|
||||
type copyFromStdinState struct {
|
||||
// copyFromStdinNode stores the original CopyFrom statement that initiated the CopyData message sequence. This
|
||||
// node is used to look at what parameters were specified, such as which table to load data into, file format,
|
||||
// delimiters, etc.
|
||||
copyFromStdinNode *node.CopyFrom
|
||||
// insertNode stores the analyzed insert node that will be used to load the data into the target table. This node
|
||||
// only needs to be built once, and can be reused with updates to its underlying data loader.
|
||||
insertNode sql.Node
|
||||
// dataLoader is the implementation of DataLoader that is used to load each individual CopyData chunk into the
|
||||
// target table.
|
||||
dataLoader dataloader.DataLoader
|
||||
// copyErr stores any error that was returned while processing a CopyData message and loading a chunk of data
|
||||
// to the target table. The server needs to keep track of any errors that were encountered while processing chunks
|
||||
// so that it can avoid sending a CommandComplete message if an error was encountered after the client already
|
||||
// sent a CopyDone message to the server.
|
||||
copyErr error
|
||||
}
|
||||
|
||||
type PortalData struct {
|
||||
Query ConvertedQuery
|
||||
IsEmptyQuery bool
|
||||
Fields []pgproto3.FieldDescription
|
||||
BoundPlan sql.Node
|
||||
FormatCodes []int16
|
||||
}
|
||||
|
||||
type PreparedStatementData struct {
|
||||
Query ConvertedQuery
|
||||
ReturnFields []pgproto3.FieldDescription
|
||||
BindVarTypes []uint32
|
||||
}
|
||||
|
||||
// extractBindVarTypes returns types based on the given query plan.
|
||||
// This function is used to get bind var types for running our prepared
|
||||
// tests only. A valid prepared query and execution messages must have
|
||||
// the types defined.
|
||||
func extractBindVarTypes(ctx *sql.Context, queryPlan sql.Node) ([]uint32, error) {
|
||||
inspectNode := queryPlan
|
||||
|
||||
types := make(map[string]uint32)
|
||||
var err error
|
||||
var extractBindVars func(ctx *sql.Context, n sql.Node, expr sql.Expression) bool
|
||||
extractBindVars = func(ctx *sql.Context, n sql.Node, expr sql.Expression) bool {
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
switch e := expr.(type) {
|
||||
// Subquery doesn't walk its Node child via Expressions, so we must walk it separately here
|
||||
case *plan.Subquery:
|
||||
transform.InspectExpressionsWithNode(ctx, e.Query, extractBindVars)
|
||||
case *expression.BindVar:
|
||||
var typOid uint32
|
||||
if doltgresType, ok := e.Type(ctx).(*pgtypes.DoltgresType); ok {
|
||||
typOid = id.Cache().ToOID(doltgresType.ID.AsId())
|
||||
} else if _, ok := e.Type(ctx).(sql.DeferredType); ok {
|
||||
// for a deferred type, we can make a guess to its type based on the containing node
|
||||
switch n.(type) {
|
||||
case *plan.Limit:
|
||||
typOid = uint32(oid.T_int4)
|
||||
case *plan.Offset:
|
||||
typOid = uint32(oid.T_int4)
|
||||
default:
|
||||
typOid, err = VitessTypeToObjectID(e.Type(ctx))
|
||||
if err != nil {
|
||||
err = errors.Errorf("could not determine OID for placeholder %s: %e", e.Name, err)
|
||||
return false
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// TODO: should remove usage non doltgres type
|
||||
typOid, err = VitessTypeToObjectID(e.Type(ctx))
|
||||
if err != nil {
|
||||
err = errors.Errorf("could not determine OID for placeholder %s: %e", e.Name, err)
|
||||
return false
|
||||
}
|
||||
}
|
||||
if existingOid, ok := types[e.Name]; ok {
|
||||
err = checkCompatibleTypes(ctx, existingOid, typOid, e.Name)
|
||||
}
|
||||
types[e.Name] = typOid
|
||||
case *pgexprs.ExplicitCast:
|
||||
if bindVar, ok := e.Child().(*expression.BindVar); ok {
|
||||
var typOid uint32
|
||||
if doltgresType, ok := bindVar.Type(ctx).(*pgtypes.DoltgresType); ok {
|
||||
typOid = id.Cache().ToOID(doltgresType.ID.AsId())
|
||||
} else {
|
||||
typOid, err = VitessTypeToObjectID(e.Type(ctx))
|
||||
if err != nil {
|
||||
err = errors.Errorf("could not determine OID for placeholder %s: %e", bindVar.Name, err)
|
||||
return false
|
||||
}
|
||||
}
|
||||
if existingOid, ok := types[bindVar.Name]; ok {
|
||||
err = checkCompatibleTypes(ctx, existingOid, typOid, bindVar.Name)
|
||||
}
|
||||
types[bindVar.Name] = typOid
|
||||
return false
|
||||
}
|
||||
// $1::text and similar get converted to a Convert expression wrapping the bindvar
|
||||
case *expression.Convert:
|
||||
if bindVar, ok := e.Child.(*expression.BindVar); ok {
|
||||
var typOid uint32
|
||||
typOid, err = VitessTypeToObjectID(e.Type(ctx))
|
||||
if err != nil {
|
||||
err = errors.Errorf("could not determine OID for placeholder %s: %e", bindVar.Name, err)
|
||||
return false
|
||||
}
|
||||
if existingOid, ok := types[bindVar.Name]; ok {
|
||||
err = checkCompatibleTypes(ctx, existingOid, typOid, bindVar.Name)
|
||||
}
|
||||
types[bindVar.Name] = typOid
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
transform.InspectExpressionsWithNode(ctx, inspectNode, extractBindVars)
|
||||
|
||||
// Insert nodes are special, as their source expressions are not returned by the Expressions()
|
||||
// interface and must be walked separately.
|
||||
switch queryPlan := queryPlan.(type) {
|
||||
case *plan.InsertInto:
|
||||
transform.InspectExpressionsWithNode(ctx, queryPlan.Source, extractBindVars)
|
||||
}
|
||||
|
||||
// above finds types of bindvars in unordered form.
|
||||
// the list of types needs to be ordered as v1, v2, v3, etc.
|
||||
typesArr := make([]uint32, len(types))
|
||||
for i, t := range types {
|
||||
idx, err := strconv.ParseInt(strings.TrimPrefix(i, "v"), 10, 32)
|
||||
if err != nil {
|
||||
return nil, errors.Errorf("could not determine the index of placeholder %s: %e", i, err)
|
||||
}
|
||||
if int(idx-1) >= len(types) {
|
||||
return nil, errors.Errorf("could not determine the index of placeholder %s in slice of %d elements", i, len(types))
|
||||
}
|
||||
typesArr[idx-1] = t
|
||||
}
|
||||
return typesArr, err
|
||||
}
|
||||
|
||||
// VitessTypeToObjectID returns a type, as defined by Vitess, into a type as defined by Postgres.
|
||||
// OIDs can be obtained with the following query: `SELECT oid, typname FROM pg_type ORDER BY 1;`
|
||||
func VitessTypeToObjectID(typ sql.Type) (uint32, error) {
|
||||
doltgresType, err := pgtypes.FromGmsTypeToDoltgresType(typ)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return id.Cache().ToOID(doltgresType.ID.AsId()), nil
|
||||
}
|
||||
|
||||
// checkCompatibleTypes checks if multiple types for which a parameter are used are compatible.
|
||||
func checkCompatibleTypes(ctx *sql.Context, existingOid, newOid uint32, newName string) error {
|
||||
var err error
|
||||
existing := pgtypes.GetTypeByID(id.Type(id.Cache().ToInternal(existingOid)))
|
||||
newType := pgtypes.GetTypeByID(id.Type(id.Cache().ToInternal(newOid)))
|
||||
if _, _, err = framework.FindCommonType(ctx, []*pgtypes.DoltgresType{existing, newType}); err != nil {
|
||||
err = errors.Errorf("parameter %s is used for incompatible types: %s and %s", newName, existing.String(), newType.String())
|
||||
}
|
||||
return err
|
||||
}
|
||||
Reference in New Issue
Block a user