dolthub--doltgresql/server/ast/table_def.go

// Copyright 2023 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 ast

import (
	"sort"

	"github.com/cockroachdb/errors"
	vitess "github.com/dolthub/vitess/go/vt/sqlparser"

	"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
	"github.com/dolthub/doltgresql/utils"
)

// assignTableDef handles tree.TableDef nodes for *vitess.DDL targets. Some table defs, such as indexes, affect other
// defs, such as columns, and they're therefore dependent on columns being handled first. It is up to the caller to
// ensure that all defs have been ordered properly before calling. assignTableDefs handles the sort for you, so this
// notice is only relevant when individually calling assignTableDef.
func assignTableDef(ctx *Context, node tree.TableDef, target *vitess.DDL) error {
	switch node := node.(type) {
	case *tree.CheckConstraintTableDef:
		if target.TableSpec == nil {
			target.TableSpec = &vitess.TableSpec{}
		}
		expr, err := nodeExpr(ctx, node.Expr)
		if err != nil {
			return err
		}
		target.TableSpec.Constraints = append(target.TableSpec.Constraints, &vitess.ConstraintDefinition{
			Name: string(node.Name),
			Details: &vitess.CheckConstraintDefinition{
				Expr:     expr,
				Enforced: true,
			},
		})
		return nil
	case *tree.ColumnTableDef:
		if target.TableSpec == nil {
			target.TableSpec = &vitess.TableSpec{}
		}
		columnDef, err := nodeColumnTableDef(ctx, node)
		if err != nil {
			return err
		}
		target.TableSpec.AddColumn(columnDef)
		if node.References.Table != nil {
			fkDef, err := nodeForeignKeyDefinitionFromColumnTableDef(ctx, node.Name, node)
			if err != nil {
				return err
			}
			target.TableSpec.Constraints = append(target.TableSpec.Constraints, &vitess.ConstraintDefinition{
				Name:    string(node.References.ConstraintName),
				Details: fkDef,
			})
		}
		return nil
	case *tree.ForeignKeyConstraintTableDef:
		if target.TableSpec == nil {
			target.TableSpec = &vitess.TableSpec{}
		}
		fkDef, err := nodeForeignKeyConstraintTableDef(ctx, node, false)
		if err != nil {
			return err
		}
		target.TableSpec.Constraints = append(target.TableSpec.Constraints, &vitess.ConstraintDefinition{
			Name:    string(node.Name),
			Details: fkDef,
		})
		return nil
	case *tree.IndexTableDef:
		if target.TableSpec == nil {
			target.TableSpec = &vitess.TableSpec{}
		}
		indexDef, err := nodeIndexTableDef(ctx, node)
		if err != nil {
			return err
		}
		target.TableSpec.Indexes = append(target.TableSpec.Indexes, indexDef)
		return nil
	case *tree.LikeTableDef:
		if len(node.Options) > 0 {
			return errors.Errorf("options for LIKE are not yet supported")
		}
		tableName, err := nodeTableName(ctx, &node.Name)
		if err != nil {
			return err
		}
		target.OptLike = &vitess.OptLike{
			LikeTables: []vitess.TableName{tableName},
		}
		return nil
	case *tree.UniqueConstraintTableDef:
		if target.TableSpec == nil {
			target.TableSpec = &vitess.TableSpec{}
		}
		indexDef, err := nodeIndexTableDef(ctx, &node.IndexTableDef)
		if err != nil {
			return err
		}
		indexDef.Info.Unique = true
		indexDef.Info.Primary = node.PrimaryKey
		// If we're setting a primary key, then we need to make sure that all of the columns are also set to NOT NULL
		if indexDef.Info.Primary {
			tableColumns := utils.SliceToMapValues(target.TableSpec.Columns, func(col *vitess.ColumnDefinition) string {
				return col.Name.String()
			})
			for _, indexedColumn := range indexDef.Fields {
				if column, ok := tableColumns[indexedColumn.Column.String()]; ok {
					column.Type.Null = false
					column.Type.NotNull = true
				}
			}
		}
		target.TableSpec.Indexes = append(target.TableSpec.Indexes, indexDef)
		return nil
	case nil:
		return nil
	default:
		return errors.Errorf("unknown table definition encountered")
	}
}

// nodeForeignKeyDefinitionFromColumnTableDef returns a vitess ForeignKeyDefinition from the specified column
// definition |node|.
func nodeForeignKeyDefinitionFromColumnTableDef(ctx *Context, fromColumn tree.Name, node *tree.ColumnTableDef) (*vitess.ForeignKeyDefinition, error) {
	if node == nil {
		return nil, nil
	}

	references := node.References
	fkConstraintTableDef := &tree.ForeignKeyConstraintTableDef{
		Name:     references.ConstraintName,
		FromCols: []tree.Name{fromColumn},
		Table:    *references.Table,
		ToCols:   []tree.Name{references.Col},
		Actions:  references.Actions,
		Match:    references.Match,
	}

	return nodeForeignKeyConstraintTableDef(ctx, fkConstraintTableDef, false)
}

// assignTableDefs handles tree.TableDefs nodes for *vitess.DDL targets. This also sorts table defs by whether they're
// dependent on other table defs evaluating first. Some table defs, such as indexes, affect other defs, such as columns,
// and they're therefore dependent on columns being handled first.
func assignTableDefs(ctx *Context, node tree.TableDefs, target *vitess.DDL) error {
	sortedNode := make(tree.TableDefs, len(node))
	copy(sortedNode, node)
	sort.Slice(sortedNode, func(i, j int) bool {
		var cmps [2]int
		for cmpsIdx := range []tree.TableDef{sortedNode[i], sortedNode[j]} {
			switch sortedNode[i].(type) {
			case *tree.IndexTableDef:
				cmps[cmpsIdx] = 1
			case *tree.UniqueConstraintTableDef:
				cmps[cmpsIdx] = 2
			default:
				cmps[cmpsIdx] = 0
			}
		}
		return cmps[0] < cmps[1]
	})
	for i := range sortedNode {
		if err := assignTableDef(ctx, sortedNode[i], target); err != nil {
			return err
		}
	}
	return nil
}