From 4d62eab6b115ae341ff2ed283f4d51dd19e61f93 Mon Sep 17 00:00:00 2001
From: yetone <yetoneful@gmail.com>
Date: Fri, 1 Jul 2022 03:44:05 +0000
Subject: [PATCH] feat: support specify pg sslmode

---
 api-server/config/yatai_config.go | 5 +++++
 api-server/services/db.go         | 9 +++++++--
 common/consts/env.go              | 1 +
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/api-server/config/yatai_config.go b/api-server/config/yatai_config.go
index 0b8eb23a..1af5ace7 100644
--- a/api-server/config/yatai_config.go
+++ b/api-server/config/yatai_config.go
@@ -22,6 +22,7 @@ type YataiPostgresqlConfigYaml struct {
 	User     string `yaml:"user"`
 	Password string `yaml:"password"`
 	Database string `yaml:"database"`
+	SSLMode  string `yaml:"sslmode"`
 }
 
 type YataiS3ConfigYaml struct {
@@ -85,6 +86,10 @@ func PopulateYataiConfig() error {
 	if ok {
 		YataiConfig.Postgresql.Database = pgDatabase
 	}
+	pgSSLMode, ok := os.LookupEnv(consts.EnvPgSSLMode)
+	if ok {
+		YataiConfig.Postgresql.SSLMode = pgSSLMode
+	}
 	migrationDir, ok := os.LookupEnv(consts.EnvMigrationDir)
 	if ok {
 		YataiConfig.Server.MigrationDir = migrationDir
diff --git a/api-server/services/db.go b/api-server/services/db.go
index d7f9b6fd..29f41229 100644
--- a/api-server/services/db.go
+++ b/api-server/services/db.go
@@ -81,12 +81,17 @@ func getPgHost() string {
 // nolint: unparam
 func getDBURI() (string, error) {
 	password := url.QueryEscape(config.YataiConfig.Postgresql.Password)
-	uri := fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=disable",
+	sslMode := "disable"
+	if config.YataiConfig.Postgresql.SSLMode != "" {
+		sslMode = config.YataiConfig.Postgresql.SSLMode
+	}
+	uri := fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=%s",
 		config.YataiConfig.Postgresql.User,
 		password,
 		getPgHost(),
 		config.YataiConfig.Postgresql.Port,
-		config.YataiConfig.Postgresql.Database)
+		config.YataiConfig.Postgresql.Database,
+		sslMode)
 	return uri, nil
 }
 
diff --git a/common/consts/env.go b/common/consts/env.go
index 7b8a945f..9805a63c 100644
--- a/common/consts/env.go
+++ b/common/consts/env.go
@@ -6,6 +6,7 @@ const (
 	EnvPgUser     = "PG_USER"
 	EnvPgPassword = "PG_PASSWORD"
 	EnvPgDatabase = "PG_DATABASE"
+	EnvPgSSLMode  = "PG_SSLMODE"
 
 	EnvMigrationDir     = "MIGRATION_DIR"
 	EnvSessionSecretKey = "SESSION_SECRET_KEY"