gormcol/gencols.go

package gormcol

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

type fieldInfo struct {
	GoName  string
	ColName string
}

type structInfo struct {
	Name     string
	Table    string
	Fields   []fieldInfo
	FilePath string
}

func parseGormColumn(tag string) string {
	for _, part := range strings.Split(tag, ";") {
		part = strings.TrimSpace(part)
		if strings.HasPrefix(part, "column:") {
			return strings.TrimPrefix(part, "column:")
		}
	}
	return ""
}

func parseModelFile(path string) (*structInfo, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
	if err != nil {
		return nil, err
	}

	var si structInfo
	si.FilePath = path

	for _, decl := range f.Decls {
		gd, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}

		if gd.Tok == token.CONST {
			for _, spec := range gd.Specs {
				vs, ok := spec.(*ast.ValueSpec)
				if !ok {
					continue
				}
				for i, name := range vs.Names {
					if strings.HasPrefix(name.Name, "TableName") {
						if i < len(vs.Values) {
							if bl, ok := vs.Values[i].(*ast.BasicLit); ok {
								si.Table = strings.Trim(bl.Value, "\"")
							}
						}
					}
				}
			}
		}

		if gd.Tok != token.TYPE {
			continue
		}
		for _, spec := range gd.Specs {
			ts, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			st, ok := ts.Type.(*ast.StructType)
			if !ok {
				continue
			}

			si.Name = ts.Name.Name

			for _, field := range st.Fields.List {
				if len(field.Names) == 0 || field.Tag == nil {
					continue
				}
				goName := field.Names[0].Name
				tag := strings.Trim(field.Tag.Value, "`")

				var gormTag string
				for _, t := range strings.Split(tag, " ") {
					t = strings.TrimSpace(t)
					if strings.HasPrefix(t, "gorm:") {
						gormTag = strings.TrimPrefix(t, "gorm:")
						gormTag = strings.Trim(gormTag, "\"")
						break
					}
				}

				colName := parseGormColumn(gormTag)
				if colName == "" {
					continue
				}

				si.Fields = append(si.Fields, fieldInfo{
					GoName:  goName,
					ColName: colName,
				})
			}
		}
	}

	if si.Name == "" {
		return nil, nil
	}

	return &si, nil
}

func generateColsVarBlock(si *structInfo) string {
	if len(si.Fields) == 0 {
		return ""
	}

	var b strings.Builder
	b.WriteString(fmt.Sprintf("\nvar %sCols = struct {\n", si.Name))
	for _, f := range si.Fields {
		b.WriteString(fmt.Sprintf("\t%s gormcol.Field\n", f.GoName))
	}
	b.WriteString("}{\n")
	for _, f := range si.Fields {
		b.WriteString(fmt.Sprintf("\t%s: gormcol.Field{Table: %q, Column: %q},\n", f.GoName, si.Table, f.ColName))
	}
	b.WriteString("}\n")
	return b.String()
}

func (m *GormGen) generateCols(fileFilter *regexp.Regexp) error {
	dir := m.cfg.OutputDir
	if !strings.HasPrefix(dir, "./") {
		dir = "./" + dir
	}

	absDir, err := filepath.Abs(dir)
	if err != nil {
		return err
	}

	entries, err := os.ReadDir(absDir)
	if err != nil {
		return err
	}

	const gormcolImport = "\"git.ma-al.com/goc_daniel/b2b/app/utils/gormcol\""

	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".go") {
			continue
		}

		// Match files by the table filter regex (strip .go suffix for matching)
		fileBase := strings.TrimSuffix(entry.Name(), ".go")
		if !fileFilter.MatchString(fileBase) {
			continue
		}

		path := filepath.Join(absDir, entry.Name())

		si, err := parseModelFile(path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "warning: skipping cols for %s: %v\n", entry.Name(), err)
			continue
		}
		if si == nil || len(si.Fields) == 0 {
			continue
		}

		content, err := os.ReadFile(path)
		if err != nil {
			return err
		}

		fileContent := string(content)

		if strings.Contains(fileContent, "gormcol.Field{") {
			continue
		}

		if !strings.Contains(fileContent, gormcolImport) {
			if strings.Contains(fileContent, "import (") {
				fileContent = strings.Replace(fileContent, "import (", "import (\n\t"+gormcolImport, 1)
			} else if strings.Contains(fileContent, "package prestadb") {
				fileContent = strings.Replace(fileContent, "package prestadb",
					"package prestadb\n\nimport "+gormcolImport, 1)
			}
		}

		colsBlock := generateColsVarBlock(si)
		fileContent += colsBlock

		if err := os.WriteFile(path, []byte(fileContent), 0644); err != nil {
			return err
		}

		fmt.Printf("Cols: %s\n", entry.Name())
	}

	return nil
}