feat: Support custom database (#910)

Signed-off-by: Ce Gao <gaoce@caicloud.io>
kubeflow · Nov 4, 2019 · d9bb39e · d9bb39e
1 parent a55cf2a
commit d9bb39e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/pkg/db/v1alpha3/common/const.go b/pkg/db/v1alpha3/common/const.go
@@ -9,7 +9,9 @@ const (
 
 	MySQLDBHostEnvName = "KATIB_MYSQL_HOST"
 	MySQLDBPortEnvName = "KATIB_MYSQL_PORT"
+	MySQLDatabase      = "KATIB_MYSQL_DATABASE"
 
-	DefaultMySQLHost = "katib-db"
-	DefaultMySQLPort = "3306"
+	DefaultMySQLDatabase = "katib"
+	DefaultMySQLHost     = "katib-db"
+	DefaultMySQLPort     = "3306"
 )
diff --git a/pkg/db/v1alpha3/mysql/mysql.go b/pkg/db/v1alpha3/mysql/mysql.go
@@ -19,7 +19,7 @@ import (
 
 const (
 	dbDriver     = "mysql"
-	dbNameTmpl   = "root:%s@tcp(%s:%s)/katib?timeout=5s"
+	dbNameTmpl   = "root:%s@tcp(%s:%s)/%s?timeout=5s"
 	mysqlTimeFmt = "2006-01-02 15:04:05.999999"
 
 	connectInterval = 5 * time.Second
@@ -37,8 +37,10 @@ func getDbName() string {
 		common.MySQLDBHostEnvName, common.DefaultMySQLHost)
 	dbPort := env.GetEnvOrDefault(
 		common.MySQLDBPortEnvName, common.DefaultMySQLPort)
+	dbName := env.GetEnvOrDefault(common.MySQLDatabase,
+		common.DefaultMySQLDatabase)
 
-	return fmt.Sprintf(dbNameTmpl, dbPass, dbHost, dbPort)
+	return fmt.Sprintf(dbNameTmpl, dbPass, dbHost, dbPort, dbName)
 }
 
 func openSQLConn(driverName string, dataSourceName string, interval time.Duration,