Add environment variable option to set postgres ssl mode

Signed-off-by: Kun Chang <curtis@mail.ustc.edu.cn>
kubeflow · Mar 1, 2024 · 3160836 · 3160836
1 parent 8df3c5c
commit 3160836
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/pkg/db/v1beta1/common/const.go b/pkg/db/v1beta1/common/const.go
@@ -41,11 +41,13 @@ const (
 	PostgreSQLDBHostEnvName = "KATIB_POSTGRESQL_DB_HOST"
 	PostgreSQLDBPortEnvName = "KATIB_POSTGRESQL_DB_PORT"
 	PostgreSQLDatabase      = "KATIB_POSTGRESQL_DB_DATABASE"
+	PostgreSslMode          = "KATIB_POSTGRESQL_SSL_MODE"
 
 	DefaultPostgreSQLUser     = "katib"
 	DefaultPostgreSQLDatabase = "katib"
 	DefaultPostgreSQLHost     = "katib-postgres"
 	DefaultPostgreSQLPort     = "5432"
+	DefaultPostgreSslMode     = "disable"
 
 	SkipDbInitializationEnvName = "SKIP_DB_INITIALIZATION"
 )
diff --git a/pkg/db/v1beta1/postgres/postgres.go b/pkg/db/v1beta1/postgres/postgres.go
@@ -48,10 +48,12 @@ func getDbName() string {
 		common.PostgreSQLDBPortEnvName, common.DefaultPostgreSQLPort)
 	dbName := env.GetEnvOrDefault(common.PostgreSQLDatabase,
 		common.DefaultPostgreSQLDatabase)
+	sslMode := env.GetEnvOrDefault(common.PostgreSslMode,
+		common.DefaultPostgreSslMode)
 
 	psqlInfo := fmt.Sprintf("host=%s port=%s user=%s "+
-		"password=%s dbname=%s sslmode=disable",
-		dbHost, dbPort, dbUser, dbPass, dbName)
+		"password=%s dbname=%s sslmode=%s",
+		dbHost, dbPort, dbUser, dbPass, dbName, sslMode)
 
 	return psqlInfo
 }

diff --git a/pkg/db/v1beta1/postgres/postgres_test.go b/pkg/db/v1beta1/postgres/postgres_test.go
@@ -137,3 +137,13 @@ func TestGetDbName(t *testing.T) {
 	}
 
 }
+
+func TestSslMode(t *testing.T) {
+	sslMode := "require"
+	t.Setenv(common.PostgreSslMode, sslMode)
+
+	dbName := "host=katib-postgres port=5432 user=katib password= dbname=katib sslmode=require"
+	if getDbName() != dbName {
+		t.Errorf("getDbName returns wrong value %v", getDbName())
+	}
+}