GoZeroPlus/tools/goctl/model/sql/parser/parser.go

package parser

import (
	"fmt"
	"strings"

	"github.com/tal-tech/go-zero/tools/goctl/model/sql/converter"
	"github.com/tal-tech/go-zero/tools/goctl/model/sql/model"
	"github.com/tal-tech/go-zero/tools/goctl/util/stringx"
	"github.com/xwb1989/sqlparser"
)

const (
	_ = iota
	primary
	unique
	normal
	spatial
)

const timeImport = "time.Time"

type (
	Table struct {
		Name       stringx.String
		PrimaryKey Primary
		Fields     []Field
	}

	Primary struct {
		Field
		AutoIncrement bool
	}

	Field struct {
		Name         stringx.String
		DataBaseType string
		DataType     string
		IsPrimaryKey bool
		IsUniqueKey  bool
		Comment      string
	}

	KeyType int
)

func Parse(ddl string) (*Table, error) {
	stmt, err := sqlparser.ParseStrictDDL(ddl)
	if err != nil {
		return nil, err
	}

	ddlStmt, ok := stmt.(*sqlparser.DDL)
	if !ok {
		return nil, errUnsupportDDL
	}

	action := ddlStmt.Action
	if action != sqlparser.CreateStr {
		return nil, fmt.Errorf("expected [CREATE] action,but found: %s", action)
	}

	tableName := ddlStmt.NewName.Name.String()
	tableSpec := ddlStmt.TableSpec
	if tableSpec == nil {
		return nil, errTableBodyNotFound
	}

	columns := tableSpec.Columns
	indexes := tableSpec.Indexes
	keyMap, err := getIndexKeyType(indexes)
	if err != nil {
		return nil, err
	}

	fields, primaryKey, err := convertFileds(columns, keyMap)
	if err != nil {
		return nil, err
	}

	return &Table{
		Name:       stringx.From(tableName),
		PrimaryKey: primaryKey,
		Fields:     fields,
	}, nil
}

func convertFileds(columns []*sqlparser.ColumnDefinition, keyMap map[string]KeyType) ([]Field, Primary, error) {
	var fields []Field
	var primaryKey Primary
	for _, column := range columns {
		if column == nil {
			continue
		}
		var comment string
		if column.Type.Comment != nil {
			comment = string(column.Type.Comment.Val)
		}
		var isDefaultNull = true
		if column.Type.NotNull {
			isDefaultNull = false
		} else {
			if column.Type.Default == nil {
				isDefaultNull = false
			} else if string(column.Type.Default.Val) != "null" {
				isDefaultNull = false
			}
		}
		dataType, err := converter.ConvertDataType(column.Type.Type, isDefaultNull)
		if err != nil {
			return nil, primaryKey, err
		}

		var field Field
		field.Name = stringx.From(column.Name.String())
		field.DataBaseType = column.Type.Type
		field.DataType = dataType
		field.Comment = comment
		key, ok := keyMap[column.Name.String()]
		if ok {
			field.IsPrimaryKey = key == primary
			field.IsUniqueKey = key == unique
			if field.IsPrimaryKey {
				primaryKey.Field = field
				if column.Type.Autoincrement {
					primaryKey.AutoIncrement = true
				}
			}
		}
		fields = append(fields, field)
	}
	return fields, primaryKey, nil
}

func getIndexKeyType(indexes []*sqlparser.IndexDefinition) (map[string]KeyType, error) {
	keyMap := make(map[string]KeyType)
	for _, index := range indexes {
		info := index.Info
		if info == nil {
			continue
		}
		if info.Primary {
			if len(index.Columns) > 1 {
				return nil, errPrimaryKey
			}

			keyMap[index.Columns[0].Column.String()] = primary
			continue
		}
		// can optimize
		if len(index.Columns) > 1 {
			continue
		}
		column := index.Columns[0]
		columnName := column.Column.String()
		camelColumnName := stringx.From(columnName).ToCamel()
		// by default, createTime|updateTime findOne is not used.
		if camelColumnName == "CreateTime" || camelColumnName == "UpdateTime" {
			continue
		}
		if info.Unique {
			keyMap[columnName] = unique
		} else if info.Spatial {
			keyMap[columnName] = spatial
		} else {
			keyMap[columnName] = normal
		}
	}
	return keyMap, nil
}

func (t *Table) ContainsTime() bool {
	for _, item := range t.Fields {
		if item.DataType == timeImport {
			return true
		}
	}
	return false
}

func ConvertColumn(db, table string, in []*model.Column) (*Table, error) {
	var reply Table
	reply.Name = stringx.From(table)
	keyMap := make(map[string][]*model.Column)

	for _, column := range in {
		keyMap[column.Key] = append(keyMap[column.Key], column)
	}
	primaryColumns := keyMap["PRI"]
	if len(primaryColumns) == 0 {
		return nil, fmt.Errorf("database:%s, table %s: missing primary key", db, table)
	}

	if len(primaryColumns) > 1 {
		return nil, fmt.Errorf("database:%s, table %s: only one primary key expected", db, table)
	}

	primaryColumn := primaryColumns[0]
	isDefaultNull := primaryColumn.ColumnDefault == nil && primaryColumn.IsNullAble == "YES"
	primaryFt, err := converter.ConvertDataType(primaryColumn.DataType, isDefaultNull)
	if err != nil {
		return nil, err
	}

	primaryField := Field{
		Name:         stringx.From(primaryColumn.Name),
		DataBaseType: primaryColumn.DataType,
		DataType:     primaryFt,
		IsUniqueKey:  true,
		IsPrimaryKey: true,
		Comment:      primaryColumn.Comment,
	}
	reply.PrimaryKey = Primary{
		Field:         primaryField,
		AutoIncrement: strings.Contains(primaryColumn.Extra, "auto_increment"),
	}
	for key, columns := range keyMap {
		for _, item := range columns {
			isColumnDefaultNull := item.ColumnDefault == nil && item.IsNullAble == "YES"
			dt, err := converter.ConvertDataType(item.DataType, isColumnDefaultNull)
			if err != nil {
				return nil, err
			}

			f := Field{
				Name:         stringx.From(item.Name),
				DataBaseType: item.DataType,
				DataType:     dt,
				IsPrimaryKey: primaryColumn.Name == item.Name,
				Comment:      item.Comment,
			}
			if key == "UNI" {
				f.IsUniqueKey = true
			}
			reply.Fields = append(reply.Fields, f)
		}
	}

	return &reply, nil
}