db/schema/dialect/duckdb/duckdb.go

package duckdb

import (
	"strings"

	"git.fsdpf.net/go/db"
	"git.fsdpf.net/go/db/dialect/duckdb"
	"git.fsdpf.net/go/db/internal/sb"
	"git.fsdpf.net/go/db/schema"
	"git.fsdpf.net/go/db/sqlgen"
)

type DuckDB struct {
	db  *db.Database
	esg sqlgen.ExpressionSQLGenerator
}

var duckdbDefaultModifiers = []string{
	"Nullable", "Default", "Increment", "Comment",
}

func (this DuckDB) GetColumnListing(table string) ([]string, error) {
	cols := []string{}
	err := this.db.From(db.S("information_schema").Table("columns")).Where(
		db.C("table_name").Eq(table),
	).Pluck(&cols, "column_name")
	return cols, err
}

func (this DuckDB) TableExists(table string) (bool, error) {
	result, err := this.db.From(db.S("information_schema").Table("tables")).Where(db.Ex{
		"table_name": table,
		"table_type": "BASE TABLE",
	}).Count()
	return result > 0, err
}

func (this DuckDB) CompileCreate(bp *schema.Blueprint) []string {
	temporary := db.L("CREATE")
	if bp.Temporary {
		temporary = db.L("CREATE TEMPORARY")
	}

	// 自增列需要先建序列
	var sqls []string
	for _, column := range bp.GetAddedColumns() {
		if column.IsAutoIncrement() {
			sqls = append(sqls, this.GenerateSQL("CREATE SEQUENCE IF NOT EXISTS ?",
				db.T(this.seqName(bp.GetTable(), column.Name))))
		}
	}

	columns := strings.Join(this.getAddedColumns(bp), ",\n")
	sqls = append(sqls, this.GenerateSQL("? TABLE IF NOT EXISTS ? (\n?\n)", temporary, db.T(bp.GetTable()), db.L(columns)))

	if comment := bp.Comment; comment != "" {
		sqls = append(sqls, this.GenerateSQL("COMMENT ON TABLE ? IS ?", db.T(bp.GetTable()), db.V(comment)))
	}
	sqls = append(sqls, this.columnComments(bp.GetTable(), bp.GetAddedColumns())...)

	return sqls
}

func (this DuckDB) CompileAdd(bp *schema.Blueprint) []string {
	columns := strings.Join(schema.PrefixArray("ADD COLUMN", this.getAddedColumns(bp)), ",\n")
	sql := this.GenerateSQL("ALTER TABLE ?\n?", db.T(bp.GetTable()), db.L(columns))
	sqls := []string{sql}
	sqls = append(sqls, this.columnComments(bp.GetTable(), bp.GetAddedColumns())...)
	return sqls
}

func (this DuckDB) CompileChange(bp *schema.Blueprint) []string {
	sqls := []string{}
	for _, column := range bp.GetChangedColumns() {
		name := column.Name
		rename := column.Name
		if column.GetRename() != "" {
			rename = column.GetRename()
		}

		// DuckDB 需要分别修改类型和约束
		if rename != name {
			sqls = append(sqls, this.GenerateSQL("ALTER TABLE ? RENAME COLUMN ? TO ?",
				db.T(bp.GetTable()), db.C(name), db.C(rename)))
		}

		// 修改类型
		sql := this.GenerateSQL("ALTER TABLE ? ALTER COLUMN ? TYPE ?",
			db.T(bp.GetTable()), db.C(rename), db.L(this.GetColumnType(column)))
		sqls = append(sqls, sql)

		// 修改 nullable
		if column.IsNullable() {
			sqls = append(sqls, this.GenerateSQL("ALTER TABLE ? ALTER COLUMN ? DROP NOT NULL",
				db.T(bp.GetTable()), db.C(rename)))
		} else {
			sqls = append(sqls, this.GenerateSQL("ALTER TABLE ? ALTER COLUMN ? SET NOT NULL",
				db.T(bp.GetTable()), db.C(rename)))
		}

		// 修改默认值
		v := column.GetDefault()
		if v != nil {
			sqls = append(sqls, this.GenerateSQL("ALTER TABLE ? ALTER COLUMN ? SET DEFAULT ?",
				db.T(bp.GetTable()), db.C(rename), v))
		} else if column.IsUseCurrent() {
			sqls = append(sqls, this.GenerateSQL("ALTER TABLE ? ALTER COLUMN ? SET DEFAULT CURRENT_TIMESTAMP",
				db.T(bp.GetTable()), db.C(rename)))
		}
	}
	return sqls
}

func (this DuckDB) CompileDropColumn(bp *schema.Blueprint) []string {
	columns := []string{}
	for _, item := range bp.GetCommands() {
		if item.Type != "dropColumn" {
			continue
		}
		for _, col := range item.Columns {
			columns = append(columns, this.GenerateSQL("DROP COLUMN ?", db.C(col)))
		}
	}
	sql := this.GenerateSQL("ALTER TABLE ?\n?", db.T(bp.GetTable()), db.L(strings.Join(columns, ",\n")))
	return []string{sql}
}

func (this DuckDB) CompileDrop(bp *schema.Blueprint) []string {
	return []string{this.GenerateSQL("DROP TABLE ?", db.T(bp.GetTable()))}
}

func (this DuckDB) CompileDropIfExists(bp *schema.Blueprint) []string {
	return []string{this.GenerateSQL("DROP TABLE IF EXISTS ?", db.T(bp.GetTable()))}
}

func (this DuckDB) CompileRename(bp *schema.Blueprint) []string {
	commands := bp.GetCommands()
	if len(commands) == 0 {
		panic("new table undefined")
	}
	toName := commands[0].To
	if toName == "" {
		panic("new table undefined")
	}
	fromName := bp.GetTable()

	sqls := []string{this.GenerateSQL("ALTER TABLE ? RENAME TO ?", db.T(fromName), db.T(toName))}

	// 查找该表关联的序列（命名规则: {table}_{column}_seq），一并重命名并更新列 DEFAULT
	var seqNames []string
	_ = this.db.From(db.L("duckdb_sequences()")).
		Where(db.C("sequence_name").Like(fromName+"_%")).
		Pluck(&seqNames, "sequence_name")

	prefix := fromName + "_"
	for _, oldSeq := range seqNames {
		if !strings.HasSuffix(oldSeq, "_seq") {
			continue
		}
		colName := oldSeq[len(prefix) : len(oldSeq)-len("_seq")]
		newSeq := this.seqName(toName, colName)
		sqls = append(sqls, this.GenerateSQL("ALTER SEQUENCE ? RENAME TO ?", db.T(oldSeq), db.T(newSeq)))
		sqls = append(sqls, this.GenerateSQL("ALTER TABLE ? ALTER COLUMN ? SET DEFAULT nextval(?)",
			db.T(toName), db.C(colName), db.V(newSeq)))
	}

	return sqls
}

func (this DuckDB) CompileModifyComment(bp *schema.Blueprint) []string {
	return []string{
		this.GenerateSQL("COMMENT ON TABLE ? IS ?", db.T(bp.GetTable()), db.V(bp.Comment)),
	}
}

func (this DuckDB) addModifiers(sql string, bp *schema.Blueprint, column *schema.ColumnDefinition) string {
	for _, modifier := range duckdbDefaultModifiers {
		sql += this.GetColumnModifier(modifier, bp, column)
	}
	return sql
}

func (this DuckDB) GetColumnType(column *schema.ColumnDefinition) string {
	switch column.Type {
	case "char":
		return this.GenerateSQL("CHAR(?)", column.Length)
	case "string":
		return this.GenerateSQL("VARCHAR(?)", column.Length)
	case "text":
		return "TEXT"
	case "integer":
		return "INTEGER"
	case "bigInteger":
		return "BIGINT"
	case "smallInteger":
		return "SMALLINT"
	case "tinyInteger":
		return "TINYINT"
	case "decimal":
		return this.GenerateSQL("DECIMAL(?,?)", column.Total, column.Places)
	case "boolean":
		return "BOOLEAN"
	case "enum":
		// DuckDB 支持 ENUM 类型
		return this.GenerateSQL("VARCHAR(?)", column.Length)
	case "json":
		return "JSON"
	case "binary":
		return "BLOB"
	case "datetime":
		return "TIMESTAMP"
	case "timestamp":
		return "TIMESTAMP"
	case "date":
		return "DATE"
	case "time":
		return "TIME"
	case "year":
		return "SMALLINT"
	case "uuid":
		return "UUID"
	case "vector":
		return this.GenerateSQL("FLOAT[?]", column.Length)
	}
	panic("Unsupported data type: " + column.Type)
}

func (this DuckDB) GetColumnModifier(modifier string, bp *schema.Blueprint, column *schema.ColumnDefinition) string {
	switch modifier {
	case "Nullable":
		if column.IsNullable() {
			return " NULL"
		}
		return " NOT NULL"
	case "Default":
		if column.IsUseCurrent() {
			// DuckDB 不支持 ON UPDATE，忽略 def 中可能携带的 MySQL ON UPDATE 标记
			return " DEFAULT CURRENT_TIMESTAMP"
		}
		v := column.GetDefault()
		if v == nil {
			return ""
		}
		return this.GenerateSQL(" DEFAULT ?", v)
	case "Increment":
		// 自增列在 getAddedColumns 中单独处理（GENERATED ALWAYS AS IDENTITY），此处无需输出
	case "Comment":
		// DuckDB 支持列注释，但需要在 CREATE TABLE 后使用 COMMENT ON
		return ""
	}
	return ""
}

func (this DuckDB) columnComments(table string, columns []*schema.ColumnDefinition) (sqls []string) {
	for _, column := range columns {
		if v := column.GetComment(); v != "" {
			sqls = append(sqls, this.GenerateSQL("COMMENT ON COLUMN ?.? IS ?",
				db.T(table), db.C(column.Name), db.V(v)))
		}
	}
	return
}

func (this DuckDB) seqName(table, column string) string {
	return table + "_" + column + "_seq"
}

func (this DuckDB) getAddedColumns(bp *schema.Blueprint) (columns []string) {
	for _, column := range bp.GetAddedColumns() {
		var sql string
		if column.IsAutoIncrement() {
			var baseType string
			switch column.Type {
			case "bigInteger":
				baseType = "BIGINT"
			case "smallInteger":
				baseType = "SMALLINT"
			default:
				baseType = "INTEGER"
			}
			seqName := this.seqName(bp.GetTable(), column.Name)
			sql = this.GenerateSQL("? ? DEFAULT nextval(?) PRIMARY KEY",
				db.C(column.Name), db.L(baseType), db.V(seqName))
		} else {
			sql = this.GenerateSQL("? ?", db.C(column.Name), db.L(this.GetColumnType(column)))
			sql = this.addModifiers(sql, bp, column)
		}
		columns = append(columns, sql)
	}
	return
}

func (this DuckDB) getChangedColumns(bp *schema.Blueprint) (columns []string) {
	for _, column := range bp.GetChangedColumns() {
		name := column.Name
		rename := column.Name
		if column.GetRename() != "" {
			rename = column.GetRename()
		}
		sql := this.GenerateSQL("? ?", db.C(rename), db.L(this.GetColumnType(column)))
		columns = append(columns, this.addModifiers(sql, bp, column))
		if rename != name {
			// 需要额外的重命名语句
			columns = append(columns, this.GenerateSQL("RENAME COLUMN ? TO ?", db.C(name), db.C(rename)))
		}
	}
	return
}

func (this DuckDB) GenerateSQL(sql string, args ...any) string {
	sb := sb.NewSQLBuilder(false)
	this.esg.Generate(sb, db.L(sql, args...))
	result, _, _ := sb.ToSQL()
	return result
}

func init() {
	schema.RegisterDialect("duckdb", func(db *db.Database) schema.Schema {
		return &DuckDB{
			db:  db,
			esg: sqlgen.NewExpressionSQLGenerator("duckdb", duckdb.DialectOptions()),
		}
	})
}