Allow for usage of env var K8S_HOST in psql (#1026)

* Allow for usage of env var `K8S_HOST` in psql

* small typo fix

* typo fix

cmd/database.go

@@ -11,36 +11,98 @@ import (
 	"gorm.io/gorm"
 )
 
-func savePgsql(jsonInfo string) {
+type PsqlConnInfo struct {
-	envVars := map[string]string{
+	Host     string
-		"PGSQL_HOST":     viper.GetString("PGSQL_HOST"),
+	User     string
-		"PGSQL_USER":     viper.GetString("PGSQL_USER"),
+	DbName   string
-		"PGSQL_DBNAME":   viper.GetString("PGSQL_DBNAME"),
+	SslMode  string
-		"PGSQL_SSLMODE":  viper.GetString("PGSQL_SSLMODE"),
+	Password string
-		"PGSQL_PASSWORD": viper.GetString("PGSQL_PASSWORD"),
+}
+
+func getPsqlConnInfo() (PsqlConnInfo, error) {
+	var host string
+	if value := viper.GetString("PGSQL_HOST"); value != "" {
+		host = value
+	} else {
+		return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_HOST env var is required", envVarsPrefix)
 	}
 
-	for k, v := range envVars {
+	var user string
-		if v == "" {
+	if value := viper.GetString("PGSQL_USER"); value != "" {
-			exitWithError(fmt.Errorf("environment variable %s is missing", envVarsPrefix+"_"+k))
+		user = value
-		}
+	} else {
+		return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_USER env var is required", envVarsPrefix)
 	}
 
-	connInfo := fmt.Sprintf("host=%s user=%s dbname=%s sslmode=%s password=%s",
+	var dbName string
-		envVars["PGSQL_HOST"],
+	if value := viper.GetString("PGSQL_DBNAME"); value != "" {
-		envVars["PGSQL_USER"],
+		dbName = value
-		envVars["PGSQL_DBNAME"],
+	} else {
-		envVars["PGSQL_SSLMODE"],
+		return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_USER env var is required", envVarsPrefix)
-		envVars["PGSQL_PASSWORD"],
+	}
+
+	var sslMode string
+	if value := viper.GetString("PGSQL_SSLMODE"); value != "" {
+		sslMode = value
+	} else {
+		return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_SSLMODE env var is required", envVarsPrefix)
+	}
+
+	var password string
+	if value := viper.GetString("PGSQL_PASSWORD"); value != "" {
+		password = value
+	} else {
+		return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_PASSWORD env var is required", envVarsPrefix)
+	}
+
+	return PsqlConnInfo{
+		Host:     host,
+		User:     user,
+		DbName:   dbName,
+		SslMode:  sslMode,
+		Password: password,
+	}, nil
+}
+
+func (c *PsqlConnInfo) toString() string {
+	return fmt.Sprintf("host=%s user=%s dbname=%s sslmode=%s password=%s",
+		c.Host,
+		c.User,
+		c.DbName,
+		c.SslMode,
+		c.Password,
 	)
+}
 
-	hostname, err := os.Hostname()
+func savePgsql(jsonInfo string) {
+	var hostname string
+	if value := viper.GetString("K8S_HOST"); value != "" {
+		// Adhere to the ScanHost column definition below
+		if len(value) > 63 {
+			exitWithError(fmt.Errorf("%s_K8S_HOST value's length must be less than 63 chars", envVarsPrefix))
+		}
+
+		hostname = value
+	} else {
+		host, err := os.Hostname()
 		if err != nil {
 			exitWithError(fmt.Errorf("received error looking up hostname: %s", err))
 		}
 
-	timestamp := time.Now()
+		hostname = host
+	}
 
+	PsqlConnInfo, err := getPsqlConnInfo()
+	if err != nil {
+		exitWithError(err)
+	}
+
+	db, err := gorm.Open(postgres.Open(PsqlConnInfo.toString()), &gorm.Config{})
+	if err != nil {
+		exitWithError(fmt.Errorf("received error connecting to database: %s", err))
+	}
+
+	timestamp := time.Now()
 	type ScanResult struct {
 		gorm.Model
 		ScanHost string    `gorm:"type:varchar(63) not null"` // https://www.ietf.org/rfc/rfc1035.txt
@@ -48,12 +110,7 @@ func savePgsql(jsonInfo string) {
 		ScanInfo string    `gorm:"type:jsonb not null"`
 	}
 
-	db, err := gorm.Open(postgres.Open(connInfo), &gorm.Config{})
-	if err != nil {
-		exitWithError(fmt.Errorf("received error connecting to database: %s", err))
-	}
-
 	db.Debug().AutoMigrate(&ScanResult{})
 	db.Save(&ScanResult{ScanHost: hostname, ScanTime: timestamp, ScanInfo: jsonInfo})
-	glog.V(2).Info(fmt.Sprintf("successfully stored result to: %s", envVars["PGSQL_HOST"]))
+	glog.V(2).Info(fmt.Sprintf("successfully stored result to: %s", PsqlConnInfo.Host))
 }