feat: support specify pg sslmode

bentoml · Jul 1, 2022 · 4d62eab · 4d62eab
1 parent 3ba8ac1
commit 4d62eab
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git 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
@@ -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
@@ -6,6 +6,7 @@ const (
 	EnvPgUser     = "PG_USER"
 	EnvPgPassword = "PG_PASSWORD"
 	EnvPgDatabase = "PG_DATABASE"
+	EnvPgSSLMode  = "PG_SSLMODE"
 
 	EnvMigrationDir     = "MIGRATION_DIR"
 	EnvSessionSecretKey = "SESSION_SECRET_KEY"