wehub-resource-sync
237 lines
9.3 KiB
Go
237 lines
9.3 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 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
|
|
}
|