package gormcol import ( "context" "fmt" "os" "path/filepath" "regexp" "strings" "gorm.io/driver/mysql" "gorm.io/gen" "gorm.io/gorm" "gorm.io/gorm/logger" ) type GenConfig struct { OutputDir string PkgName string TableFilter string // regex pattern, e.g. "(ps_|b2b_).*" } func defaultConfig() GenConfig { return GenConfig{ OutputDir: "./app/model/dbmodel", PkgName: "dbmodel", TableFilter: "ps_.*", } } type GormGen struct { db *gorm.DB cfg GenConfig } func New(db *gorm.DB) *GormGen { return &GormGen{db: db, cfg: defaultConfig()} } func NewWithConfig(db *gorm.DB, cfg GenConfig) *GormGen { d := defaultConfig() if cfg.OutputDir != "" { d.OutputDir = cfg.OutputDir } if cfg.PkgName != "" { d.PkgName = cfg.PkgName } if cfg.TableFilter != "" { d.TableFilter = cfg.TableFilter } return &GormGen{db: db, cfg: d} } // ConnectDSN opens a MySQL/MariaDB connection from a DSN string. func ConnectDSN(dsn string) (*gorm.DB, error) { db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{ Logger: logger.Default.LogMode(logger.Error), }) if err != nil { return nil, fmt.Errorf("failed to connect with dsn: %w", err) } return db, nil } func (m *GormGen) GenModels(ctx context.Context) error { re, err := regexp.Compile("^" + m.cfg.TableFilter + "$") if err != nil { return fmt.Errorf("invalid table filter regex %q: %w", m.cfg.TableFilter, err) } g := gen.NewGenerator(gen.Config{ OutPath: m.cfg.OutputDir, ModelPkgPath: m.cfg.PkgName, FieldNullable: true, FieldWithIndexTag: true, }) g.UseDB(m.db) tableNames, err := m.db.Migrator().GetTables() if err != nil { return fmt.Errorf("failed to get table list: %w", err) } matched := 0 for _, tableName := range tableNames { if re.MatchString(tableName) { g.GenerateModel(tableName) matched++ } } fmt.Printf("Matched %d tables with filter %q\n", matched, m.cfg.TableFilter) g.Execute() if err := m.cleanupGeneratedFiles(); err != nil { return fmt.Errorf("failed to cleanup generated files: %w", err) } if err := m.generateCols(re); err != nil { return fmt.Errorf("failed to generate column descriptors: %w", err) } return nil } func (m *GormGen) cleanupGeneratedFiles() error { filesToRemove := []string{"gen.go", "do.go", "_gen.go"} dir := m.cfg.OutputDir if !strings.HasPrefix(dir, "./") { dir = "./" + dir } absDir, err := filepath.Abs(dir) if err != nil { return err } for _, fileName := range filesToRemove { filePath := filepath.Join(absDir, fileName) if _, err := os.Stat(filePath); err == nil { if err := os.Remove(filePath); err != nil { return fmt.Errorf("failed to remove %s: %w", filePath, err) } fmt.Printf("Removed: %s\n", filePath) } } files, err := os.ReadDir(absDir) if err != nil { return err } re, err := regexp.Compile("^(" + m.cfg.TableFilter + ")\\.gen\\.go$") if err != nil { return err } for _, file := range files { name := file.Name() if re.MatchString(name) { oldPath := filepath.Join(absDir, name) baseName := strings.TrimSuffix(name, ".gen.go") newPath := filepath.Join(absDir, baseName+".go") content, err := os.ReadFile(oldPath) if err != nil { return err } content = m.cleanModelContent(content) if err := os.WriteFile(newPath, content, 0644); err != nil { return err } if err := os.Remove(oldPath); err != nil { return err } fmt.Printf("Renamed: %s -> %s\n", oldPath, newPath) } } return nil } func (m *GormGen) cleanModelContent(content []byte) []byte { result := string(content) lines := strings.Split(result, "\n") var newLines []string importStarted := false importEnded := false for _, line := range lines { trimmed := strings.TrimSpace(line) if trimmed == "import (" { importStarted = true newLines = append(newLines, line) continue } if importStarted && trimmed == ")" { importEnded = true importStarted = false newLines = append(newLines, line) continue } if importStarted && !importEnded { if strings.Contains(trimmed, "\"gorm.io/gen\"") || strings.Contains(trimmed, "\"gorm.io/gen/field\"") || strings.Contains(trimmed, "\"gorm.io/plugin/dbresolver\"") || strings.Contains(trimmed, "gen.DO") { continue } } if strings.Contains(trimmed, "type psCategoryDo struct") || strings.HasPrefix(trimmed, "func (p psCategoryDo)") { continue } newLines = append(newLines, line) } result = strings.Join(newLines, "\n") result = strings.ReplaceAll(result, "psCategoryDo", "") return []byte(result) }