From 202239795bf96cd0089919dba9f388f4f4056e7a Mon Sep 17 00:00:00 2001
From: Rogee <rogeecn@qq.com>
Date: Fri, 12 Sep 2025 14:15:28 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E7=A7=BB=E9=99=A4=E5=86=97=E4=BD=99?=
 =?UTF-8?q?=E7=9A=84=E6=95=B0=E6=8D=AE=E5=BA=93=E9=85=8D=E7=BD=AE=E5=A4=84?=
 =?UTF-8?q?=E7=90=86=EF=BC=8C=E7=AE=80=E5=8C=96=E6=A8=A1=E5=9E=8B=E7=94=9F?=
 =?UTF-8?q?=E6=88=90=E5=91=BD=E4=BB=A4=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/gen_model.go | 108 ++++++++++++-----------------------------------
 1 file changed, 27 insertions(+), 81 deletions(-)

diff --git a/cmd/gen_model.go b/cmd/gen_model.go
index c9aa2e8..e54d23f 100644
--- a/cmd/gen_model.go
+++ b/cmd/gen_model.go
@@ -1,16 +1,14 @@
 package cmd
 
 import (
-	"fmt"
-
-	"github.com/pkg/errors"
-	log "github.com/sirupsen/logrus"
-	"github.com/spf13/cobra"
-	"github.com/spf13/viper"
-	"go.ipao.vip/atomctl/v2/pkg/utils/gomod"
-	"go.ipao.vip/gen"
-	"gorm.io/driver/postgres"
-	"gorm.io/gorm"
+    "github.com/pkg/errors"
+    log "github.com/sirupsen/logrus"
+    "github.com/spf13/cobra"
+    apg "go.ipao.vip/atomctl/v2/pkg/postgres"
+    "go.ipao.vip/atomctl/v2/pkg/utils/gomod"
+    "go.ipao.vip/gen"
+    "gorm.io/driver/postgres"
+    "gorm.io/gorm"
 )
 
 func CommandGenModel(root *cobra.Command) {
@@ -38,83 +36,31 @@ func CommandGenModel(root *cobra.Command) {
 }
 
 func commandGenModelE(cmd *cobra.Command, args []string) error {
-	if err := gomod.Parse("go.mod"); err != nil {
-		return errors.Wrap(err, "parse go.mod")
-	}
+    if err := gomod.Parse("go.mod"); err != nil {
+        return errors.Wrap(err, "parse go.mod")
+    }
 
-	cfgFile := cmd.Flag("config").Value.String()
-	if cfgFile == "" {
-		cfgFile = "config.toml"
-	}
+    cfgFile := cmd.Flag("config").Value.String()
+    if cfgFile == "" {
+        cfgFile = "config.toml"
+    }
 
-	v := viper.New()
-	v.SetConfigType("toml")
-	v.SetConfigFile(cfgFile)
-	v.AddConfigPath(".")
-	if err := v.ReadInConfig(); err != nil {
-		return errors.Wrap(err, "read config")
-	}
+    sqlDB, conf, err := apg.GetDB(cfgFile)
+    if err != nil {
+        return errors.Wrap(err, "load database config")
+    }
+    defer sqlDB.Close()
 
-	var dbc modelDBConfig
-	if err := v.UnmarshalKey("Database", &dbc); err != nil {
-		return errors.Wrap(err, "unmarshal Database config")
-	}
-	dsn := dbc.DSN()
+    dsn := conf.DSN()
+    log.Infof("parsed DSN: %s (schema=%s)", dsn, conf.Schema)
 
-	log.Infof("parsed DSN: %s (schema=%s)", dsn, dbc.Schema)
-
-	db, err := gorm.Open(postgres.New(postgres.Config{DSN: dsn}))
-	if err != nil {
-		return errors.Wrapf(err, "open database with dsn: %s", dsn)
-	}
+    db, err := gorm.Open(postgres.New(postgres.Config{DSN: dsn}))
+    if err != nil {
+        return errors.Wrapf(err, "open database with dsn: %s", dsn)
+    }
 
 	// 默认同包同目录生成到 ./database
 	gen.GenerateWithDefault(db, "./database/.transform.yaml")
 
-	return nil
-}
-
-// local config and helpers
-
-type modelDBConfig struct {
-	Username string
-	Password string
-	Database string
-	Schema   string
-	Host     string
-	Port     uint
-	SslMode  string
-	TimeZone string
-}
-
-func (c *modelDBConfig) applyDefaults() {
-	if c.Username == "" {
-		c.Username = "postgres"
-	}
-	if c.SslMode == "" {
-		c.SslMode = "disable"
-	}
-	if c.TimeZone == "" {
-		c.TimeZone = "Asia/Shanghai"
-	}
-	if c.Port == 0 {
-		c.Port = 5432
-	}
-	if c.Schema == "" {
-		c.Schema = "public"
-	}
-}
-
-func (c *modelDBConfig) DSN() string {
-	c.applyDefaults()
-	return fmt.Sprintf(
-		"host=%s user=%s password=%s dbname=%s port=%d sslmode=%s TimeZone=%s",
-		c.Host,
-		c.Username,
-		c.Password,
-		c.Database,
-		c.Port,
-		c.SslMode,
-		c.TimeZone,
-	)
+    return nil
 }