gormcol/gen.go

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