@@ -15,12 +15,14 @@ import (
1515 "strings"
1616 "sync"
1717
18+ "github.com/go-sql-driver/mysql"
1819 "github.com/jzero-io/jzero-contrib/filex"
1920 "github.com/pkg/errors"
2021 "github.com/samber/lo"
2122 "github.com/spf13/cast"
2223 "github.com/zeromicro/go-zero/core/color"
2324 "github.com/zeromicro/go-zero/core/logx"
25+ "github.com/zeromicro/go-zero/core/stores/postgres"
2426 "github.com/zeromicro/go-zero/core/stores/sqlx"
2527 "github.com/zeromicro/go-zero/tools/goctl/model/sql/parser"
2628 "github.com/zeromicro/go-zero/tools/goctl/util/format"
@@ -66,14 +68,29 @@ func (jm *JzeroModel) Gen() error {
6668 sqlConn sqlx.SqlConn
6769 )
6870
69- if config .C .Gen .ModelMysqlDatasource {
71+ if config .C .Gen .ModelDriver == "postgres" && ! config .C .Gen .ModelDatasource {
72+ return errors .New ("postgres model only support datasource mode" )
73+ }
74+
75+ if config .C .Gen .ModelMysqlDatasource || config .C .Gen .ModelDatasource {
7076 if jm .IsNew {
7177 fmt .Printf ("%s you are using mysql datesource to generate model code, please manual execute jzero gen command\n " , color .WithColor ("Detected" , color .FgRed ))
7278 return nil
7379 }
74- sqlConn = sqlx .NewMysql (config .C .Gen .ModelMysqlDatasourceUrl )
7580
76- tables , err := getMysqlAllTables (sqlConn )
81+ switch config .C .Gen .ModelDriver {
82+ case "mysql" :
83+ if config .C .Gen .ModelMysqlDatasourceUrl != "" {
84+ config .C .Gen .ModelDatasourceUrl = config .C .Gen .ModelMysqlDatasourceUrl
85+ }
86+ sqlConn = sqlx .NewMysql (config .C .Gen .ModelDatasourceUrl )
87+ case "postgres" :
88+ sqlConn = postgres .New (config .C .Gen .ModelDatasourceUrl )
89+ default :
90+ return errors .Errorf ("model driver %s not support" , config .C .Gen .ModelDriver )
91+ }
92+
93+ tables , err := getAllTables (sqlConn , config .C .Gen .ModelDriver )
7794 if err != nil {
7895 return err
7996 }
@@ -84,7 +101,7 @@ func (jm *JzeroModel) Gen() error {
84101 if err != nil {
85102 return err
86103 }
87- if ! config .C .Gen .MysqlCreateTableDDL {
104+ if ! config .C .Gen .ModelMysqlCreateTableDDL || ! config . C . Gen . ModelCreateTableDDL {
88105 defer func () {
89106 for _ , v := range writeTables {
90107 if err = os .Remove (v ); err != nil {
@@ -161,18 +178,18 @@ func (jm *JzeroModel) Gen() error {
161178 var mu sync.Mutex
162179
163180 if len (genCodeSqlFiles ) != 0 {
164- if config .C .Gen .ModelMysqlDatasource {
165- tables , err := getMysqlAllTables (sqlConn )
181+ if config .C .Gen .ModelMysqlDatasource || config . C . Gen . ModelDatasource {
182+ tables , err := getAllTables (sqlConn , config . C . Gen . ModelDriver )
166183 if err != nil {
167184 return err
168185 }
169- if len (config .C .Gen .ModelMysqlDatasourceTable ) != 0 && config .C .Gen .ModelMysqlDatasourceTable [0 ] != "*" {
186+ if ( len (config .C .Gen .ModelMysqlDatasourceTable ) != 0 && config .C .Gen .ModelMysqlDatasourceTable [0 ] != "*" ) || ( len ( config . C . Gen . ModelDatasourceTable ) != 0 && config . C . Gen . ModelDatasourceTable [ 0 ] != "*" ) {
170187 for _ , v := range tables {
171- if lo .Contains (config .C .Gen .ModelMysqlDatasourceTable , cast .ToString (v )) {
188+ if lo .Contains (config .C .Gen .ModelMysqlDatasourceTable , cast .ToString (v )) || lo . Contains ( config . C . Gen . ModelDatasourceTable , cast . ToString ( v )) {
172189 allTables = append (allTables , v )
173190 }
174191 }
175- } else if len (config .C .Gen .ModelMysqlDatasourceTable ) != 0 && config .C .Gen .ModelMysqlDatasourceTable [0 ] == "*" {
192+ } else if ( len (config .C .Gen .ModelMysqlDatasourceTable ) != 0 && config .C .Gen .ModelMysqlDatasourceTable [0 ] == "*" ) || ( len ( config . C . Gen . ModelDatasourceTable ) != 0 && config . C . Gen . ModelDatasourceTable [ 0 ] == "*" ) {
176193 allTables = tables
177194 }
178195 for _ , f := range allFiles {
@@ -218,12 +235,38 @@ func (jm *JzeroModel) Gen() error {
218235
219236 bf := filepath .Base (f )
220237 modelDir := filepath .Join ("internal" , "model" , strings .ToLower (bf [0 :len (bf )- len (path .Ext (bf ))]))
221- cmd := exec .Command ("goctl" , "model" , "mysql" , "ddl" , "--database" , config .C .Gen .ModelMysqlDDLDatabase , "--src" , f , "--dir" , modelDir , "--home" , goctlHome , "--style" , config .C .Gen .Style , "-i" , strings .Join (config .C .Gen .ModelMysqlIgnoreColumns , "," ), "--cache=" + fmt .Sprintf ("%t" , config .C .Gen .ModelMysqlCache ), "--strict=" + fmt .Sprintf ("%t" , config .C .Gen .ModelMysqlStrict ))
222- resp , err := cmd .CombinedOutput ()
223- if err != nil {
224- return errors .Errorf ("gen model code meet error. Err: %s:%s" , err .Error (), resp )
238+
239+ var ddlDatabase string
240+ if config .C .Gen .ModelMysqlDDLDatabase != "" {
241+ ddlDatabase = config .C .Gen .ModelMysqlDDLDatabase
242+ } else if config .C .Gen .ModelMysqlDatasourceUrl != "" {
243+ mysqlDsn , err := mysql .ParseDSN (config .C .Gen .ModelMysqlDatasourceUrl )
244+ if err != nil {
245+ return err
246+ }
247+ ddlDatabase = mysqlDsn .DBName
225248 }
226- if config .C .Gen .ModelMysqlCachePrefix != "" && config .C .Gen .ModelMysqlCache {
249+
250+ if config .C .Gen .ModelDriver == "postgres" {
251+ cmd := exec .Command ("goctl" , "model" , "pg" , "datasource" , "--url" , config .C .Gen .ModelDatasourceUrl , "-t" , strings .TrimSuffix (filepath .Base (f ), ".sql" ), "--dir" , modelDir , "--home" , goctlHome , "--style" , config .C .Gen .Style , "-i" , strings .Join (config .C .Gen .ModelIgnoreColumns , "," ), "--cache=" + fmt .Sprintf ("%t" , config .C .Gen .ModelCache ), "--strict=" + fmt .Sprintf ("%t" , config .C .Gen .ModelStrict ))
252+ logx .Debug (cmd .String ())
253+ resp , err := cmd .CombinedOutput ()
254+ if err != nil {
255+ return errors .Errorf ("gen model code meet error. Err: %s:%s" , err .Error (), resp )
256+ }
257+ } else {
258+ if config .C .Gen .ModelMysqlCache {
259+ config .C .Gen .ModelCache = true
260+ }
261+ cmd := exec .Command ("goctl" , "model" , "mysql" , "ddl" , "--database" , ddlDatabase , "--src" , f , "--dir" , modelDir , "--home" , goctlHome , "--style" , config .C .Gen .Style , "-i" , strings .Join (config .C .Gen .ModelMysqlIgnoreColumns , "," ), "--cache=" + fmt .Sprintf ("%t" , config .C .Gen .ModelCache ), "--strict=" + fmt .Sprintf ("%t" , config .C .Gen .ModelMysqlStrict ))
262+ logx .Debug (cmd .String ())
263+ resp , err := cmd .CombinedOutput ()
264+ if err != nil {
265+ return errors .Errorf ("gen model code meet error. Err: %s:%s" , err .Error (), resp )
266+ }
267+ }
268+
269+ if (config .C .Gen .ModelMysqlCachePrefix != "" && config .C .Gen .ModelMysqlCache ) || (config .C .Gen .ModelCachePrefix != "" && config .C .Gen .ModelCache ) {
227270 for _ , tp := range tableParsers {
228271 namingFormat , err := format .FileNamingFormat (config .C .Gen .Style , tp .Name .Source ())
229272 if err != nil {
@@ -233,7 +276,7 @@ func (jm *JzeroModel) Gen() error {
233276 if config .C .Gen .Style == "go_zero" {
234277 file = namingFormat + "_model_gen.go"
235278 }
236- err = jm .addModelMysqlCachePrefix (filepath .Join (modelDir , file ))
279+ err = jm .addModelCachePrefix (filepath .Join (modelDir , file ))
237280 if err != nil {
238281 return err
239282 }
@@ -251,7 +294,7 @@ func (jm *JzeroModel) Gen() error {
251294 return nil
252295}
253296
254- func (jm * JzeroModel ) addModelMysqlCachePrefix (fp string ) error {
297+ func (jm * JzeroModel ) addModelCachePrefix (fp string ) error {
255298 fset := token .NewFileSet ()
256299 f , err := goparser .ParseFile (fset , fp , nil , goparser .ParseComments )
257300 if err != nil {
@@ -266,7 +309,11 @@ func (jm *JzeroModel) addModelMysqlCachePrefix(fp string) error {
266309 if strings .HasPrefix (name .Name , "cache" ) && strings .HasSuffix (name .Name , "Prefix" ) {
267310 value := valueSpec .Values [i ]
268311 if basicLit , ok := value .(* ast.BasicLit ); ok {
269- basicLit .Value = fmt .Sprintf (`"%s%s"` , config .C .Gen .ModelMysqlCachePrefix , strings .ReplaceAll (basicLit .Value , "\" " , "" ))
312+ if config .C .Gen .ModelCachePrefix != "" {
313+ basicLit .Value = fmt .Sprintf (`"%s%s"` , config .C .Gen .ModelCachePrefix , strings .ReplaceAll (basicLit .Value , "\" " , "" ))
314+ } else if config .C .Gen .ModelMysqlCachePrefix != "" {
315+ basicLit .Value = fmt .Sprintf (`"%s%s"` , config .C .Gen .ModelMysqlCachePrefix , strings .ReplaceAll (basicLit .Value , "\" " , "" ))
316+ }
270317 }
271318 }
272319 }
@@ -286,15 +333,20 @@ func (jm *JzeroModel) addModelMysqlCachePrefix(fp string) error {
286333 return nil
287334}
288335
289- type Table struct {
290- Name string `db:"name"`
291- }
292-
293- func getMysqlAllTables (sqlConn sqlx.SqlConn ) ([]string , error ) {
336+ func getAllTables (sqlConn sqlx.SqlConn , driver string ) ([]string , error ) {
294337 var tables []string
295- err := sqlConn .QueryRowsCtx (context .Background (), & tables , "show tables" )
296- if err != nil {
297- return nil , err
338+
339+ switch driver {
340+ case "mysql" :
341+ err := sqlConn .QueryRowsCtx (context .Background (), & tables , "show tables" )
342+ if err != nil {
343+ return nil , err
344+ }
345+ case "postgres" :
346+ err := sqlConn .QueryRowsCtx (context .Background (), & tables , "select tablename from pg_tables where schemaname = 'public'" )
347+ if err != nil {
348+ return nil , err
349+ }
298350 }
299351 return tables , nil
300352}
@@ -303,7 +355,11 @@ type ShowCreateTableResult struct {
303355 DDL string `db:"Create Table"`
304356}
305357
306- func getTableDDL (sqlConn sqlx.SqlConn , table string ) (string , error ) {
358+ func getTableDDL (sqlConn sqlx.SqlConn , driver , table string ) (string , error ) {
359+ if driver == "postgres" {
360+ return "-- todo" , nil
361+ }
362+
307363 var showCreateTableResult ShowCreateTableResult
308364 err := sqlConn .QueryRowCtx (context .Background (), & showCreateTableResult , "show create table " + table )
309365 if err != nil {
0 commit comments