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") } columns := strings.Join(this.getAddedColumns(bp), ",\n") sql := this.GenerateSQL("? TABLE ? (\n?\n)", temporary, db.T(bp.GetTable()), db.L(columns)) if comment := bp.Comment; comment != "" { sql = sql[:len(sql)-1] + ";" + this.GenerateSQL("\nCOMMENT ON TABLE ? IS ?", db.T(bp.GetTable()), db.V(comment)) } return []string{sql} } 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)) return []string{sql} } 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 { toName := "" commands := bp.GetCommands() if len(commands) == 0 { panic("new table undefined") } toName = commands[0].To if toName == "" { panic("new table undefined") } return []string{this.GenerateSQL("ALTER TABLE ? RENAME TO ?", db.T(bp.GetTable()), db.T(toName))} } 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" } 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": v := column.GetDefault() if v == nil { if column.IsUseCurrent() { return " DEFAULT CURRENT_TIMESTAMP" } return "" } return this.GenerateSQL(" DEFAULT ?", v) case "Increment": if column.IsAutoIncrement() { // DuckDB 使用 SERIAL 或者 SEQUENCE return " PRIMARY KEY" } case "Comment": // DuckDB 支持列注释,但需要在 CREATE TABLE 后使用 COMMENT ON return "" } return "" } func (this DuckDB) getAddedColumns(bp *schema.Blueprint) (columns []string) { for _, column := range bp.GetAddedColumns() { colType := this.GetColumnType(column) // 对于自增列,使用 SERIAL 类型 if column.IsAutoIncrement() { switch column.Type { case "bigInteger": colType = "BIGSERIAL" case "smallInteger": colType = "SMALLSERIAL" default: colType = "SERIAL" } } sql := this.GenerateSQL("? ?", db.C(column.Name), db.L(colType)) columns = append(columns, this.addModifiers(sql, bp, column)) } 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()), } }) }