#90: add connection params to psql init

IBM · Sep 20, 2017 · a7f0f51 · a7f0f51
1 parent 6c33a0e
commit a7f0f51
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 10 deletions.
diff --git a/database/init.go b/database/init.go
@@ -28,11 +28,15 @@ const (
     KeyPsqlPassword = "UBIQUITY_DB_PASSWORD"
     KeyPsqlDbName = "UBIQUITY_DB_NAME"
     KeyPsqlPort = "UBIQUITY_DB_PORT"
+    KeyPsqlTimeout = "UBIQUITY_DB_CONNECT_TIMEOUT"
+    KeyPsqlSslMode = "UBIQUITY_DB_SSL_MODE"
+    KeyPsqlSslRootCert = "UBIQUITY_DB_SSL_ROOT_CRT"
 )
 
 func GetPsqlConnectionParams(hostname string) string {
+    str := ""
     // add host
-    str := "host=" + hostname
+    str += "host=" + hostname
     // add user
     psqlUser := os.Getenv(KeyPsqlUser)
     if psqlUser == "" {
@@ -55,11 +59,28 @@ func GetPsqlConnectionParams(hostname string) string {
     if psqlPort != "" {
         str += " port=" + psqlPort
     }
+    // add connect_timeout
+    psqlTimeout := os.Getenv(KeyPsqlTimeout)
+    if psqlTimeout != "" {
+        str += " connect_timeout=" + psqlTimeout
+    }
     return str
 }
 
 func GetPsqlSslParams() string {
-    return "sslmode=disable"
+    str := ""
+    // add sslmode
+    psqlSslMode := os.Getenv(KeyPsqlSslMode)
+    if psqlSslMode == "" {
+        psqlSslMode = "disable"
+    }
+    str += "sslmode=" + psqlSslMode
+    // add sslrootcert
+    psqlSslRootCert := os.Getenv(KeyPsqlSslRootCert)
+    if psqlSslRootCert != "" {
+        str += " sslrootcert=" + psqlSslRootCert
+    }
+    return str
 }
 
 func InitPostgres(hostname string) func() {

diff --git a/database/init_test.go b/database/init_test.go
@@ -28,18 +28,22 @@ import (
 
 var _ = Describe("Init", func() {
     var (
-        hostname      string = "my-hostname"
-        defaultUser   string = "postgres"
-        defaultDbName string = "postgres"
-        newUser       string = "my-user"
-        newDbName     string = "my-dbname"
-        newPassword   string = "my-password"
-        newPort       string = "my-port"
+        hostname       string = "my-hostname"
+        defaultUser    string = "postgres"
+        defaultDbName  string = "postgres"
+        newUser        string = "my-user"
+        newDbName      string = "my-dbname"
+        newPassword    string = "my-password"
+        newPort        string = "my-port"
+        newTimeout     string = "my-timeout"
+        defaultSslMode string = "disable"
+        newSslMode     string = "my-sslmode"
+        newSslRootCert string = "my-sslrootcert"
     )
     BeforeEach(func() {
     })
 
-    Context(".Postgres", func() {
+    Context(".Postgres connection", func() {
         It("only hostname", func() {
             res := database.GetPsqlConnectionParams(hostname)
             Expect(res).To(Equal(fmt.Sprintf("host=%s user=%s dbname=%s", hostname, defaultUser, defaultDbName)))
@@ -80,5 +84,31 @@ var _ = Describe("Init", func() {
             os.Unsetenv(database.KeyPsqlPort)
             Expect(res).To(Equal(fmt.Sprintf("host=%s user=%s dbname=%s password=%s port=%s", hostname, newUser, newDbName, newPassword, newPort)))
         })
+        It("hostname timeout", func() {
+            os.Setenv(database.KeyPsqlTimeout, newTimeout)
+            res := database.GetPsqlConnectionParams(hostname)
+            os.Unsetenv(database.KeyPsqlTimeout)
+            Expect(res).To(Equal(fmt.Sprintf("host=%s user=%s dbname=%s connect_timeout=%s", hostname, defaultUser, defaultDbName, newTimeout)))
+        })
+    })
+    Context(".Postgres ssl", func() {
+        It("empty", func() {
+            res := database.GetPsqlSslParams()
+            Expect(res).To(Equal(fmt.Sprintf("sslmode=%s", defaultSslMode)))
+        })
+        It("sslmode", func() {
+            os.Setenv(database.KeyPsqlSslMode, newSslMode)
+            res := database.GetPsqlSslParams()
+            os.Unsetenv(database.KeyPsqlSslMode)
+            Expect(res).To(Equal(fmt.Sprintf("sslmode=%s", newSslMode)))
+        })
+        It("sslmode sslrootcert", func() {
+            os.Setenv(database.KeyPsqlSslMode, newSslMode)
+            os.Setenv(database.KeyPsqlSslRootCert, newSslRootCert)
+            res := database.GetPsqlSslParams()
+            os.Unsetenv(database.KeyPsqlSslMode)
+            os.Unsetenv(database.KeyPsqlSslRootCert)
+            Expect(res).To(Equal(fmt.Sprintf("sslmode=%s sslrootcert=%s", newSslMode, newSslRootCert)))
+        })
     })
 })