dolthub--doltgresql/server/analyzer/assign_insert_casts.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 analyzer

import (
	"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"
	"github.com/dolthub/go-mysql-server/sql/plan"
	"github.com/dolthub/go-mysql-server/sql/transform"
	"github.com/dolthub/go-mysql-server/sql/types"

	pgexprs "github.com/dolthub/doltgresql/server/expression"
	"github.com/dolthub/doltgresql/server/node"
	pgtypes "github.com/dolthub/doltgresql/server/types"
)

// AssignInsertCasts adds the appropriate assign casts for insertions.
func AssignInsertCasts(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope *plan.Scope, selector analyzer.RuleSelector, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
	insertInto, ok := node.(*plan.InsertInto)
	if !ok {
		return node, transform.SameTree, nil
	}

	// We have some sources that are already postgres native, so skip them
	if isDoltgresNativeSource(ctx, insertInto.Destination, insertInto.Source) {
		return insertInto, transform.SameTree, nil
	}

	// First we'll make a map for each column, so that it's easier to match a name to a type. We also ensure that the
	// types use Doltgres types, as casts rely on them. At this point, we shouldn't have any GMS types floating around
	// anymore, so no need to include a lot of additional code to handle them.
	destinationNameToType := make(map[string]*pgtypes.DoltgresType)
	for _, col := range insertInto.Destination.Schema(ctx) {
		colType, ok := col.Type.(*pgtypes.DoltgresType)
		if !ok {
			// Only non-Doltgres destination tables will have GMS types (such as system tables), so we don't error here
			colType = pgtypes.FromGmsType(col.Type)
		}
		destinationNameToType[strings.ToLower(col.Name)] = colType
	}

	// Create the destination type slice that will match each inserted column
	destinationTypes := make([]*pgtypes.DoltgresType, len(insertInto.ColumnNames))
	for i, colName := range insertInto.ColumnNames {
		destinationTypes[i], ok = destinationNameToType[strings.ToLower(colName)]
		if !ok {
			return nil, transform.NewTree, errors.Errorf("INSERT: cannot find destination column with name `%s`", colName)
		}
	}

	// Replace expressions with casts as needed
	if values, ok := insertInto.Source.(*plan.Values); ok {
		// Values do not return the correct Schema since each row may contain different types, so we must handle it differently
		newValues := make([][]sql.Expression, len(values.ExpressionTuples))
		for rowIndex, rowExprs := range values.ExpressionTuples {
			newValues[rowIndex] = make([]sql.Expression, len(rowExprs))
			for columnIndex, colExpr := range rowExprs {
				// Null ColumnDefaultValues or empty DefaultValues are not properly typed in TypeSanitizer, so we must handle them here
				colExprType := colExpr.Type(ctx)
				if colExprType == nil || colExprType == types.Null {
					colExprType = pgtypes.Unknown
				}
				fromColType, ok := colExprType.(*pgtypes.DoltgresType)
				if !ok {
					return nil, transform.NewTree, errors.Errorf("INSERT: non-Doltgres type found in values source: %s", fromColType.String())
				}
				toColType := destinationTypes[columnIndex]
				// We only assign the existing expression if the types perfectly match (same parameters), otherwise we'll cast
				if fromColType.Equals(toColType) {
					newValues[rowIndex][columnIndex] = colExpr
				} else {
					newValues[rowIndex][columnIndex] = pgexprs.NewAssignmentCast(colExpr, fromColType, toColType)
				}
			}
		}
		insertInto = insertInto.WithSource(plan.NewValues(newValues))
	} else {
		sourceSchema := insertInto.Source.Schema(ctx)
		projections := make([]sql.Expression, len(sourceSchema))
		for i, col := range sourceSchema {
			fromColType, ok := col.Type.(*pgtypes.DoltgresType)
			if !ok {
				return nil, transform.NewTree, errors.Errorf("INSERT: non-Doltgres type found in source: %s", fromColType.String())
			}
			toColType := destinationTypes[i]
			getField := expression.NewGetField(i, fromColType, col.Name, true)
			// We only assign the GetField if the types perfectly match (same parameters), otherwise we'll cast
			if fromColType.Equals(toColType) {
				projections[i] = getField
			} else {
				projections[i] = pgexprs.NewAssignmentCast(getField, fromColType, toColType)
			}
		}
		insertInto = insertInto.WithSource(plan.NewProject(ctx, projections, insertInto.Source))
	}

	// handle on conflict clause if present
	if insertInto.OnDupExprs.HasUpdates() {
		newDupExprs, err := assignUpdateFieldCasts(ctx, insertInto.OnDupExprs.AllExpressions())
		if err != nil {
			return nil, false, err
		}
		// TODO: this relies on a particular implementation detail InsertInto.WithExpressions
		newInsertInto, err := insertInto.WithExpressions(ctx, append(newDupExprs, insertInto.Checks().ToExpressions()...)...)
		if err != nil {
			return nil, false, err
		}

		insertInto = newInsertInto.(*plan.InsertInto)
	}

	return insertInto, transform.NewTree, nil
}

func isDoltgresNativeSource(ctx *sql.Context, dest sql.Node, source sql.Node) bool {
	// we still need this transformation if our destination has generated columns
	for _, col := range dest.Schema(ctx) {
		if col.Generated != nil {
			return false
		}
	}

	switch source.(type) {
	case *node.CopyFrom:
		return true
	default:
		return false
	}
}