From bdee945461775bcc48d3e96d06ac72b046d328d1 Mon Sep 17 00:00:00 2001
From: rogee <rogeecn@gmail.com>
Date: Wed, 1 Mar 2023 09:03:40 +0800
Subject: [PATCH] support pgsql model generation

---
 cmd/model.go                         | 23 ++++++++++++++----
 config.toml                          | 15 +++++++++++-
 providers/config/section_database.go | 22 ++++++++++-------
 providers/database/database.go       |  2 ++
 providers/database/postgres.go       | 36 ++++++++++++++++++++++++++++
 5 files changed, 84 insertions(+), 14 deletions(-)
 create mode 100644 providers/database/postgres.go

diff --git a/cmd/model.go b/cmd/model.go
index 47aab05..34453a1 100644
--- a/cmd/model.go
+++ b/cmd/model.go
@@ -5,6 +5,8 @@ package cmd
 
 import (
 	"atom/container"
+	"atom/providers/config"
+	"atom/providers/database"
 	"errors"
 	"log"
 
@@ -18,7 +20,8 @@ import (
 type GenQueryGenerator struct {
 	dig.In
 
-	DB *gorm.DB
+	Config *config.Config
+	DB     *gorm.DB
 }
 
 // // Dynamic SQL
@@ -35,10 +38,22 @@ var modelCmd = &cobra.Command{
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return container.Container.Invoke(func(gq GenQueryGenerator) error {
 			var tables []string
-			err := gq.DB.Raw("show tables").Scan(&tables).Error
-			if err != nil {
-				log.Fatal(err)
+
+			switch gq.Config.Database.Driver {
+			case database.DriverMySQL:
+				err := gq.DB.Raw("show tables").Scan(&tables).Error
+				if err != nil {
+					log.Fatal(err)
+				}
+			case database.DriverPostgres:
+				err := gq.DB.Raw("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'").Scan(&tables).Error
+				if err != nil {
+					log.Fatal(err)
+				}
+			case database.DriverSQLite:
+
 			}
+
 			if len(tables) == 0 {
 				return errors.New("no tables in database, run migrate first")
 			}
diff --git a/config.toml b/config.toml
index b709ace..e597f03 100644
--- a/config.toml
+++ b/config.toml
@@ -88,7 +88,7 @@ ShowLine = true
 LogInConsole = true
 
 [Database]
-Driver = "mysql"
+Driver = "postgres"
 
 [Database.MySQL]
 Host = "10.47.119.226"
@@ -102,6 +102,19 @@ MaxIdleConns = 10
 MaxOpenConns = 200
 Engine = "InnoDB"
 
+[Database.PostgreSQL]
+Host = "10.95.31.212"
+Port = 5442
+Database = "cspm"
+Username = "postgres"
+Password = "zHTz7H6S9WJQK7Y0"
+Prefix = ""
+TimeZone = "Asia/Shanghai"
+SslMode = "disable"
+Singular = false
+MaxIdleConns = 10
+MaxOpenConns = 200
+
 [Database.SQLite]
 File = "sqlite.db"
 
diff --git a/providers/config/section_database.go b/providers/config/section_database.go
index 093ea0f..accbaf9 100755
--- a/providers/config/section_database.go
+++ b/providers/config/section_database.go
@@ -58,25 +58,29 @@ func (m *MySQL) DSN() string {
 }
 
 type PostgreSQL struct {
-	User     string
-	Password string
-	Database string
-	Host     string
-	Port     uint
-	SslMode  string
-	TimeZone string
+	Username     string
+	Password     string
+	Database     string
+	Host         string
+	Port         uint
+	SslMode      string
+	TimeZone     string
+	Prefix       string // 表前缀
+	Singular     bool   //是否开启全局禁用复数，true表示开启
+	MaxIdleConns int    // 空闲中的最大连接数
+	MaxOpenConns int    // 打开到数据库的最大连接数
 }
 
 func (m *PostgreSQL) EmptyDsn() string {
 	dsnTpl := "host=%s user=%s password=%s port=%d dbname=postgres sslmode=disable TimeZone=Asia/Shanghai"
 
-	return fmt.Sprintf(dsnTpl, m.Host, m.User, m.Password, m.Port)
+	return fmt.Sprintf(dsnTpl, m.Host, m.Username, m.Password, m.Port)
 }
 
 // DSN connection dsn
 func (m *PostgreSQL) DSN() string {
 	dsnTpl := "host=%s user=%s password=%s dbname=%s port=%d sslmode=%s TimeZone=%s"
-	return fmt.Sprintf(dsnTpl, m.Host, m.User, m.Password, m.Database, m.Port, m.SslMode, m.TimeZone)
+	return fmt.Sprintf(dsnTpl, m.Host, m.Username, m.Password, m.Database, m.Port, m.SslMode, m.TimeZone)
 }
 
 type Redis struct {
diff --git a/providers/database/database.go b/providers/database/database.go
index 5b286cc..33d80ba 100644
--- a/providers/database/database.go
+++ b/providers/database/database.go
@@ -28,6 +28,8 @@ func NewDatabase(config *config.Config) (*gorm.DB, error) {
 		return NewMySQL(config.Database.MySQL)
 	case DriverSQLite:
 		return NewSQLite(config.Database.SQLite)
+	case DriverPostgres:
+		return NewPostgres(config.Database.PostgreSQL)
 	}
 	return nil, errors.New("failed to connect to db")
 }
diff --git a/providers/database/postgres.go b/providers/database/postgres.go
new file mode 100644
index 0000000..464cf0e
--- /dev/null
+++ b/providers/database/postgres.go
@@ -0,0 +1,36 @@
+package database
+
+import (
+	"atom/providers/config"
+	"log"
+
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+	"gorm.io/gorm/schema"
+)
+
+func NewPostgres(conf *config.PostgreSQL) (*gorm.DB, error) {
+	dbConfig := postgres.Config{
+		DSN: conf.DSN(), // DSN data source name
+	}
+	log.Println("PostgreSQL DSN: ", dbConfig.DSN)
+
+	gormConfig := gorm.Config{
+		NamingStrategy: schema.NamingStrategy{
+			TablePrefix:   conf.Prefix,
+			SingularTable: conf.Singular,
+		},
+		DisableForeignKeyConstraintWhenMigrating: true,
+	}
+
+	db, err := gorm.Open(postgres.New(dbConfig), &gormConfig)
+	if err != nil {
+		return nil, err
+	}
+
+	sqlDB, _ := db.DB()
+	sqlDB.SetMaxIdleConns(conf.MaxIdleConns)
+	sqlDB.SetMaxOpenConns(conf.MaxOpenConns)
+
+	return db, err
+}