Merge pull request showwin#4 from luvtechno/load_env

Load database credentials from env vars
Goryudyuma · May 2, 2018 · c1a9408 · c1a9408
2 parents 72f593f + 7417e0d
commit c1a9408
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 7 deletions.
diff --git a/admin/benchmark.go b/admin/benchmark.go
@@ -1,9 +1,11 @@
 package main
 
 import (
+	"database/sql"
 	"flag"
 	"fmt"
 	"log"
+	"os"
 	"strconv"
 	"sync"
 	"time"
@@ -29,6 +31,21 @@ func loopBakugaiScenario(wg *sync.WaitGroup, m *sync.Mutex, finishTime time.Time
 	}
 }
 
+func getEnv(key, fallback string) string {
+	if value, ok := os.LookupEnv(key); ok {
+		return value
+	}
+	return fallback
+}
+
+func getDB() (*sql.DB, error) {
+	user := getEnv("ISHOCON1_DB_USER", "ishocon")
+	pass := getEnv("ISHOCON1_DB_PASSWORD", "ishocon")
+	dbname := getEnv("ISHOCON1_DB_NAME", "ishocon1")
+	db, err := sql.Open("mysql", user+":"+pass+"@/"+dbname)
+	return db, err
+}
+
 func startBenchmark(workload int) {
 	getInitialize()
 	log.Print("Benchmark Start!  Workload: " + strconv.Itoa(workload))

diff --git a/admin/support.go b/admin/support.go
@@ -1,7 +1,6 @@
 package main
 
 import (
-	"database/sql"
 	"math/rand"
 	"time"
 
@@ -20,7 +19,7 @@ func getUserInfo(id int) (int, string, string) {
 		id = getRand(1, 5000)
 	}
 	var email, password string
-	db, err := sql.Open("mysql", "ishocon:ishocon@/ishocon1")
+	db, err := getDB()
 	if err != nil {
 		panic(err.Error())
 	}

diff --git a/admin/validator.go b/admin/validator.go
@@ -1,7 +1,6 @@
 package main
 
 import (
-	"database/sql"
 	"log"
 	"net/http"
 	"os"
@@ -224,7 +223,7 @@ func validateUsers(id int, loggedIn bool) {
 }
 
 func getTotalPay(userID int) string {
-	db, err := sql.Open("mysql", "ishocon:ishocon@/ishocon1")
+	db, err := getDB()
 	if err != nil {
 		panic(err.Error())
 	}

diff --git a/webapp/go/main.go b/webapp/go/main.go
@@ -16,11 +16,18 @@ import (
 
 var db *sql.DB
 
+func getEnv(key, fallback string) string {
+	if value, ok := os.LookupEnv(key); ok {
+		return value
+	}
+	return fallback
+}
+
 func main() {
 	// database setting
-	user := os.Getenv("ISHOCON1_DB_USER")
-	pass := os.Getenv("ISHOCON1_DB_PASSWORD")
-	dbname := "ishocon1"
+	user := getEnv("ISHOCON1_DB_USER", "ishocon")
+	pass := getEnv("ISHOCON1_DB_PASSWORD", "ishocon")
+	dbname := getEnv("ISHOCON1_DB_NAME", "ishocon1")
 	db, _ = sql.Open("mysql", user+":"+pass+"@/"+dbname)
 	db.SetMaxIdleConns(5)