From 647a15d4bc34138331c37e569f654b24bc0f88dc Mon Sep 17 00:00:00 2001 From: quentinkhoo Date: Tue, 21 Jan 2025 15:33:06 +0800 Subject: [PATCH 1/4] add configurable options for ping check and sslmode --- README.md | 2 ++ kshieldconfig_example.toml | 1 + pkg/config/config.go | 1 + pkg/mysqldb/connect.go | 11 +++++++---- pkg/postgresdb/connect.go | 22 ++++++++++++---------- 5 files changed, 23 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index d2ff29f..cef0dc5 100644 --- a/README.md +++ b/README.md @@ -100,11 +100,13 @@ port="5432" user="postgres" dbname="postgres" password="xxxxx" +sslmode="disable" maxIdleConn = 2 maxOpenConn = 2 [app] debug = true +pingCheck = true ``` diff --git a/kshieldconfig_example.toml b/kshieldconfig_example.toml index beab5e9..840a03e 100644 --- a/kshieldconfig_example.toml +++ b/kshieldconfig_example.toml @@ -5,6 +5,7 @@ # port = "5432" # user = "postgres" # password = "password123" +# sslmode = "disable" # dbname = "mydb" # maxIdleConn = 10 # maxOpenConn = 100 diff --git a/pkg/config/config.go b/pkg/config/config.go index 9cf3182..91d8b7f 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -235,6 +235,7 @@ type GeneratePassword struct { type App struct { Debug bool `toml:"debug"` Hostname string `toml:"hostname"` + PingCheck bool `toml:"pingCheck"` Run bool RunPostgres bool RunMySql bool diff --git a/pkg/mysqldb/connect.go b/pkg/mysqldb/connect.go index 5aea3e2..7d3f80c 100644 --- a/pkg/mysqldb/connect.go +++ b/pkg/mysqldb/connect.go @@ -28,10 +28,13 @@ func Open(conf config.MySQL) (*sql.DB, string, error) { Msg("Failed to connect to database") return nil, "", err } - err = db.Ping() - if err != nil { - // fmt.Printf("Failed to connect to database. Error: %s", err.Error()) - return nil, "", err + + if cnf.App.PingCheck { + err = db.Ping() + if err != nil { + // fmt.Printf("Failed to connect to database. Error: %s", err.Error()) + return nil, "", err + } } if conf.MaxIdleConn > 0 { db.SetMaxIdleConns(conf.MaxIdleConn) diff --git a/pkg/postgresdb/connect.go b/pkg/postgresdb/connect.go index 71f0a34..08e19c0 100644 --- a/pkg/postgresdb/connect.go +++ b/pkg/postgresdb/connect.go @@ -15,7 +15,7 @@ type Postgres struct { User string `toml:"user"` Password string `toml:"password"` DBName string `toml:"dbname"` - // SSLmode string `toml:"sslmode"` + SSLmode string `toml:"sslmode"` MaxIdleConn int `toml:"maxIdleConn"` MaxOpenConn int `toml:"maxOpenConn"` } @@ -35,7 +35,7 @@ var re = regexp.MustCompile(`(?m)(?:host=)([^\s]+)`) func Open(conf Postgres) (*sql.DB, string, error) { // "host=localhost port=5432 user=postgres password=postgres dbname=postgres sslmode=disable" - url := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable", conf.Host, conf.Port, conf.User, conf.Password, conf.DBName) + url := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s", conf.Host, conf.Port, conf.User, conf.Password, conf.DBName) db, err := ConnectDatabaseUsingConnectionString(url) if err != nil { @@ -78,14 +78,16 @@ func ConnectDatabaseUsingConnectionString(url string) (*sql.DB, error) { return nil, err } - err = db.Ping() - if err != nil { - log.Error(). - Err(err). - Str("conn", url). - Msg("Failed to ping database") - db.Close() - return nil, err + if cnf.App.PingCheck { + err = db.Ping() + if err != nil { + log.Error(). + Err(err). + Str("conn", url). + Msg("Failed to ping database") + db.Close() + return nil, err + } } return db, nil From 97307faec5230f29e75991681e64de78230542d0 Mon Sep 17 00:00:00 2001 From: quentinkhoo Date: Tue, 21 Jan 2025 15:53:43 +0800 Subject: [PATCH 2/4] refactor --- README.md | 1 - kshieldconfig_example.toml | 1 + pkg/config/config.go | 1 - pkg/mysqldb/connect.go | 2 +- pkg/postgresdb/connect.go | 3 ++- 5 files changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index cef0dc5..c07f2a6 100644 --- a/README.md +++ b/README.md @@ -106,7 +106,6 @@ maxOpenConn = 2 [app] debug = true -pingCheck = true ``` diff --git a/kshieldconfig_example.toml b/kshieldconfig_example.toml index 840a03e..bc46d1b 100644 --- a/kshieldconfig_example.toml +++ b/kshieldconfig_example.toml @@ -6,6 +6,7 @@ # user = "postgres" # password = "password123" # sslmode = "disable" +# pingcheck = true # dbname = "mydb" # maxIdleConn = 10 # maxOpenConn = 100 diff --git a/pkg/config/config.go b/pkg/config/config.go index 91d8b7f..9cf3182 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -235,7 +235,6 @@ type GeneratePassword struct { type App struct { Debug bool `toml:"debug"` Hostname string `toml:"hostname"` - PingCheck bool `toml:"pingCheck"` Run bool RunPostgres bool RunMySql bool diff --git a/pkg/mysqldb/connect.go b/pkg/mysqldb/connect.go index 7d3f80c..1375007 100644 --- a/pkg/mysqldb/connect.go +++ b/pkg/mysqldb/connect.go @@ -29,7 +29,7 @@ func Open(conf config.MySQL) (*sql.DB, string, error) { return nil, "", err } - if cnf.App.PingCheck { + if conf.PingCheck { err = db.Ping() if err != nil { // fmt.Printf("Failed to connect to database. Error: %s", err.Error()) diff --git a/pkg/postgresdb/connect.go b/pkg/postgresdb/connect.go index 08e19c0..17844cf 100644 --- a/pkg/postgresdb/connect.go +++ b/pkg/postgresdb/connect.go @@ -16,6 +16,7 @@ type Postgres struct { Password string `toml:"password"` DBName string `toml:"dbname"` SSLmode string `toml:"sslmode"` + PingCheck bool `toml:"pingCheck"` MaxIdleConn int `toml:"maxIdleConn"` MaxOpenConn int `toml:"maxOpenConn"` } @@ -78,7 +79,7 @@ func ConnectDatabaseUsingConnectionString(url string) (*sql.DB, error) { return nil, err } - if cnf.App.PingCheck { + if conf.PingCheck { err = db.Ping() if err != nil { log.Error(). From ec07837daf39deade2f9e263bce08065c5f3beeb Mon Sep 17 00:00:00 2001 From: quentinkhoo Date: Tue, 21 Jan 2025 17:12:29 +0800 Subject: [PATCH 3/4] try --- pkg/config/config.go | 1 + pkg/postgresdb/connect.go | 10 ++++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/pkg/config/config.go b/pkg/config/config.go index 9cf3182..74a9543 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -215,6 +215,7 @@ type MySQL struct { Port string `toml:"port"` User string `toml:"user"` Password string `toml:"password"` + PingCheck bool `toml:"pingCheck"` // DBName string `toml:"dbname"` // SSLmode string `toml:"sslmode"` MaxIdleConn int `toml:"maxIdleConn"` diff --git a/pkg/postgresdb/connect.go b/pkg/postgresdb/connect.go index 17844cf..e402bec 100644 --- a/pkg/postgresdb/connect.go +++ b/pkg/postgresdb/connect.go @@ -4,6 +4,7 @@ import ( "database/sql" "fmt" "regexp" + "strings" _ "github.com/lib/pq" "github.com/rs/zerolog/log" @@ -36,7 +37,7 @@ var re = regexp.MustCompile(`(?m)(?:host=)([^\s]+)`) func Open(conf Postgres) (*sql.DB, string, error) { // "host=localhost port=5432 user=postgres password=postgres dbname=postgres sslmode=disable" - url := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s", conf.Host, conf.Port, conf.User, conf.Password, conf.DBName) + url := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s pingcheck=%t", conf.Host, conf.Port, conf.User, conf.Password, conf.DBName, conf.SSLmode, conf.PingCheck) db, err := ConnectDatabaseUsingConnectionString(url) if err != nil { @@ -69,7 +70,11 @@ func Open(conf Postgres) (*sql.DB, string, error) { // ConnectDatabaseUsingConnectionString connects to a PostgreSQL database using the provided connection string. // It returns a database connection, the connection string, and an error if any. + func ConnectDatabaseUsingConnectionString(url string) (*sql.DB, error) { + parts := strings.Split(url, "pingcheck=") + pingcheck := parts[1] + url = parts[0] db, err := sql.Open("postgres", url) if err != nil { log.Error(). @@ -79,7 +84,7 @@ func ConnectDatabaseUsingConnectionString(url string) (*sql.DB, error) { return nil, err } - if conf.PingCheck { + if len(pingcheck) > 0 && pingcheck == "true" { err = db.Ping() if err != nil { log.Error(). @@ -91,5 +96,6 @@ func ConnectDatabaseUsingConnectionString(url string) (*sql.DB, error) { } } + return db, nil } From 852d16719d8958002e051663bdc40767d535194b Mon Sep 17 00:00:00 2001 From: quentinkhoo Date: Tue, 21 Jan 2025 17:24:02 +0800 Subject: [PATCH 4/4] Update kshieldconfig --- README.md | 1 + kshieldconfig_example.toml | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index c07f2a6..641cb0a 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,7 @@ user="postgres" dbname="postgres" password="xxxxx" sslmode="disable" +pingCheck=true maxIdleConn = 2 maxOpenConn = 2 diff --git a/kshieldconfig_example.toml b/kshieldconfig_example.toml index bc46d1b..9a60492 100644 --- a/kshieldconfig_example.toml +++ b/kshieldconfig_example.toml @@ -6,7 +6,7 @@ # user = "postgres" # password = "password123" # sslmode = "disable" -# pingcheck = true +# pingCheck = true # dbname = "mydb" # maxIdleConn = 10 # maxOpenConn = 100