initial

gen.go

@@ -0,0 +1,214 @@
+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)
+}