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()), } }) }