gormcol/cmd/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"strings"
	"time"

	"git.ma-al.com/goc_marek/gormcol"
	"github.com/joho/godotenv"
	"gorm.io/gorm"
)

// main is the entry point for the gormcol-gen CLI tool.
// It parses flags, loads configuration from .env, connects to the database,
// and generates GORM models with column descriptors.
func main() {
	dsn := flag.String("dsn", "", "database DSN (e.g. user:pass@tcp(host:3306)/dbname)")
	filter := flag.String("filter", "(ps_|b2b_).*", "regex to match table names")
	all := flag.Bool("all", false, "generate all tables matching filter (shows confirmation)")
	outDir := flag.String("out", "./app/model/dbmodel", "output directory for generated files")
	pkgName := flag.String("pkg", "dbmodel", "Go package name for generated files")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "gormcol-gen - generate GORM models with column descriptors\n\n")
		fmt.Fprintf(os.Stderr, "Usage:\n")
		fmt.Fprintf(os.Stderr, "  gormcol-gen [options]\n\n")
		fmt.Fprintf(os.Stderr, "DSN can be provided via --dsn flag or DSN env var (from .env)\n\n")
		fmt.Fprintf(os.Stderr, "Flags:\n")
		flag.PrintDefaults()
	}

	// Load .env file if present in current directory.
	godotenv.Load()

	flag.Parse()

	// Get DSN from flag or environment variable.
	dsnValue := *dsn
	if dsnValue == "" {
		dsnValue = os.Getenv("DSN")
	}
	if dsnValue == "" {
		flag.Usage()
		fmt.Fprintln(os.Stderr, "\nerror: --dsn or DSN env is required")
		os.Exit(1)
	}

	// Connect to the database.
	db, err := gormcol.ConnectDSN(dsnValue)
	if err != nil {
		slog.Error("failed to connect to database", "error", err)
		os.Exit(1)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	cfg := gormcol.GenConfig{
		OutputDir:   *outDir,
		PkgName:     *pkgName,
		TableFilter: *filter,
	}

	// Interactive mode (default): select tables using fzf.
	// All mode: generate all tables matching filter with confirmation.
	if *all {
		if !confirmGenerateAll(db) {
			fmt.Println("Aborted.")
			os.Exit(0)
		}
	} else {
		// Check if fzf is available for interactive mode.
		if _, err := exec.LookPath("fzf"); err != nil {
			fmt.Fprintln(os.Stderr, "error: fzf is required for interactive mode. Install from https://github.com/junegunn/fzf")
			os.Exit(1)
		}
		selected, err := selectTablesInteractive(db)
		if err != nil {
			slog.Error("failed to select tables interactively", "error", err)
			os.Exit(1)
		}
		if len(selected) == 0 {
			fmt.Println("No tables selected, exiting.")
			os.Exit(0)
		}
		cfg.SelectedTables = selected
	}

	// Generate models with column descriptors.
	if err := gormcol.NewWithConfig(db, cfg).GenModels(ctx); err != nil {
		slog.Error("failed to generate models", "error", err)
		os.Exit(1)
	}
}

// selectTablesInteractive displays all database tables in fzf for interactive selection.
// Returns a list of selected table names.
func selectTablesInteractive(db *gorm.DB) ([]string, error) {
	type migratorWithGetTables interface {
		GetTables() ([]string, error)
	}

	tableNames, err := db.Migrator().(migratorWithGetTables).GetTables()
	if err != nil {
		return nil, fmt.Errorf("failed to get tables: %w", err)
	}

	fmt.Printf("Found %d tables, launching fzf for selection...\n", len(tableNames))

	selected, err := runFzf(tableNames)
	if err != nil {
		return nil, fmt.Errorf("fzf selection failed: %w", err)
	}

	return selected, nil
}

// confirmGenerateAll displays a confirmation prompt before generating all tables.
// Returns true if user confirms, false if cancelled or timeout.
func confirmGenerateAll(db *gorm.DB) bool {
	type migratorWithGetTables interface {
		GetTables() ([]string, error)
	}

	tableNames, err := db.Migrator().(migratorWithGetTables).GetTables()
	if err != nil {
		return false
	}

	if len(tableNames) == 0 {
		fmt.Println("No tables in database.")
		return false
	}

	// Display yellow warning message and wait for user input.
	yellow := "\033[33m"
	reset := "\033[0m"
	msg := fmt.Sprintf("%sWARNING:%s Generate all %d tables? [Enter] confirm / [Esc] cancel", yellow, reset, len(tableNames))

	script := fmt.Sprintf(`echo; echo -e '%s'; read -t 10 -n 1 -s key; if [[ -z "$key" ]]; then exit 0; elif [[ "$key" == $'\e' ]]; then exit 1; else exit 0; fi`, msg)
	cmd := exec.Command("bash", "-c", script)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err = cmd.Run()
	return err == nil
}

// runFzf launches fzf for interactive table selection.
// Supports multi-select with Tab, fuzzy search, and keyboard navigation.
func runFzf(tableNames []string) ([]string, error) {
	header := "Tab: toggle select | Enter: confirm | Esc: cancel"
	cmd := exec.Command("fzf", "-m", "--height=50%", "--border", "--prompt", "Select tables> ", "--header", header)
	cmd.Stdin = strings.NewReader(strings.Join(tableNames, "\n"))

	out, err := cmd.Output()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			if exitErr.ExitCode() == 130 {
				// User pressed Esc to cancel.
				return []string{}, nil
			}
		}
		return nil, fmt.Errorf("fzf failed: %w", err)
	}

	// Parse selected lines from fzf output.
	selected := strings.Split(strings.TrimSpace(string(out)), "\n")
	var result []string
	for _, s := range selected {
		s = strings.TrimSpace(s)
		if s != "" {
			result = append(result, s)
		}
	}
	return result, nil
}