go-db/sqldef/generators/doris.go

package generators

import (
	"fmt"
	"strings"
)

// DorisGenerator Doris SQL生成器
type DorisGenerator struct{}

// NewDorisGenerator 创建Doris生成器实例
func NewDorisGenerator() *DorisGenerator {
	return &DorisGenerator{}
}

func (dg *DorisGenerator) DBType() string {
	return "doris"
}

func (dg *DorisGenerator) TableExistsSQL(tableName string) string {
	return fmt.Sprintf(
		"SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = '%s'",
		tableName,
	)
}

func (dg *DorisGenerator) DropTableSQL(tableName string) string {
	return fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName)
}

func (dg *DorisGenerator) GenerateCreateTable(table TableDDL) string {
	if table.Schema == nil {
		return ""
	}

	var sql strings.Builder
	var primaryKeyColumns []string

	// 表头
	sql.WriteString(fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (\n", table.Name))

	// 列定义 - 第一遍收集主键列
	for _, col := range table.Schema.Columns {
		for _, opt := range col.Options {
			if strings.ToUpper(opt) == "PRIMARY KEY" {
				primaryKeyColumns = append(primaryKeyColumns, col.Name)
				break
			}
		}
	}

	// 列定义 - 第二遍生成列定义
	for i, col := range table.Schema.Columns {
		colDef := fmt.Sprintf("  %s %s", col.Name, dg.getDorisType(col))

		// 处理列选项，过滤掉 PRIMARY KEY（不在列级定义）
		for _, opt := range col.Options {
			upperOpt := strings.ToUpper(opt)
			if upperOpt != "PRIMARY KEY" && upperOpt != "UNIQUE" {
				colDef += " " + opt
			}
		}

		// 如果是主键列，确保有 NOT NULL
		isPrimaryKey := false
		for _, pk := range primaryKeyColumns {
			if pk == col.Name {
				isPrimaryKey = true
				break
			}
		}

		if isPrimaryKey {
			hasNotNull := false
			for _, opt := range col.Options {
				if strings.ToUpper(opt) == "NOT NULL" {
					hasNotNull = true
					break
				}
			}
			if !hasNotNull {
				colDef += " NOT NULL"
			}
		}

		// 添加默认值
		if col.Default != "" {
			upperDefault := strings.ToUpper(col.Default)
			if strings.Contains(upperDefault, "CURRENT_TIMESTAMP") ||
				strings.Contains(upperDefault, "NOW()") ||
				strings.Contains(upperDefault, "UUID()") ||
				strings.Contains(upperDefault, "NULL") {
				colDef += " DEFAULT " + col.Default
			} else {
				colDef += fmt.Sprintf(" DEFAULT '%s'", col.Default)
			}
		}

		// 添加注释
		if col.Comment != "" {
			colDef += fmt.Sprintf(` COMMENT "%s"`, col.Comment)
		}

		if i < len(table.Schema.Columns)-1 {
			sql.WriteString(colDef + ",\n")
		} else {
			sql.WriteString(colDef + "\n")
		}
	}

	// 关闭列定义括号
	sql.WriteString(")")

	// 添加 UNIQUE KEY（如果有主键列）
	if len(primaryKeyColumns) > 0 {
		sql.WriteString(fmt.Sprintf("\nUNIQUE KEY(%s)", strings.Join(primaryKeyColumns, ", ")))
	}

	// 分布式配置 - 使用第一个主键列，如果没有主键则使用第一个列
	distKey := ""
	if len(primaryKeyColumns) > 0 {
		distKey = primaryKeyColumns[0]
	} else if len(table.Schema.Columns) > 0 {
		distKey = table.Schema.Columns[0].Name
	}

	if distKey != "" {
		sql.WriteString(fmt.Sprintf("\nDISTRIBUTED BY HASH(%s) BUCKETS 10", distKey))
	}

	// 副本数配置
	sql.WriteString("\nPROPERTIES (\"replication_num\" = \"1\");")

	return sql.String()
}

// getDorisType 获取Doris数据类型
func (dg *DorisGenerator) getDorisType(col ColumnSchema) string {
	switch col.Type {
	case "DECIMAL":
		if col.Precision > 0 && col.Scale > 0 {
			return fmt.Sprintf("DECIMAL(%d,%d)", col.Precision, col.Scale)
		}
		return "DECIMAL"
	case "VARCHAR":
		if col.Length > 0 {
			return fmt.Sprintf("VARCHAR(%d)", col.Length)
		}
		return "VARCHAR(255)"
	case "CHAR":
		if col.Length > 0 {
			return fmt.Sprintf("CHAR(%d)", col.Length)
		}
		return "CHAR(1)"
	case "INT":
		return "INT"
	case "BIGINT":
		return "BIGINT"
	case "TINYINT":
		return "TINYINT"
	case "BOOL":
		return "BOOLEAN"
	case "DATETIME":
		return "DATETIME"
	case "TIMESTAMP":
		return "DATETIME" // Doris 中通常用 DATETIME
	case "DATE":
		return "DATE"
	case "TIME":
		return "TIME"
	case "TEXT":
		return "STRING"
	case "JSON":
		return "STRING"
	case "BLOB":
		return "STRING"
	case "FLOAT":
		return "FLOAT"
	case "DOUBLE":
		return "DOUBLE"
	default:
		return col.Type
	}
}

func init() {
	RegisterGenerator(NewDorisGenerator())
}