Add --connect-timeout flag to katib-db-manager (#1937)

kubeflow · Aug 22, 2022 · ca903e2 · ca903e2
1 parent 9c7d797
commit ca903e2
Show file tree

Hide file tree

Showing 7 changed files with 25 additions and 13 deletions.
diff --git a/cmd/db-manager/v1beta1/Dockerfile b/cmd/db-manager/v1beta1/Dockerfile
@@ -28,4 +28,3 @@ WORKDIR /app
 COPY --from=build-env /bin/grpc_health_probe /bin/
 COPY --from=build-env /go/src/github.com/kubeflow/katib/katib-db-manager /app/
 ENTRYPOINT ["./katib-db-manager"]
-CMD ["-w", "kubernetes"]
diff --git a/cmd/db-manager/v1beta1/main.go b/cmd/db-manager/v1beta1/main.go
@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"time"
 
 	health_pb "github.com/kubeflow/katib/pkg/apis/manager/health"
 	api_pb "github.com/kubeflow/katib/pkg/apis/manager/v1beta1"
@@ -34,7 +35,8 @@ import (
 )
 
 const (
-	port = "0.0.0.0:6789"
+	port                  = "0.0.0.0:6789"
+	defaultConnectTimeout = time.Second * 60
 )
 
 var dbIf common.KatibDBInterface
@@ -87,14 +89,17 @@ func (s *server) Check(ctx context.Context, in *health_pb.HealthCheckRequest) (*
 }
 
 func main() {
+	var connectTimeout time.Duration
+	flag.DurationVar(&connectTimeout, "connect-timeout", defaultConnectTimeout, "Timeout before calling error during database connection. (e.g. 120s)")
 	flag.Parse()
+
 	var err error
 	dbNameEnvName := common.DBNameEnvName
 	dbName := os.Getenv(dbNameEnvName)
 	if dbName == "" {
 		klog.Fatal("DB_NAME env is not set. Exiting")
 	}
-	dbIf, err = db.NewKatibDBInterface(dbName)
+	dbIf, err = db.NewKatibDBInterface(dbName, connectTimeout)
 	if err != nil {
 		klog.Fatalf("Failed to open db connection: %v", err)
 	}

diff --git a/docs/developer-guide.md b/docs/developer-guide.md
@@ -58,7 +58,7 @@ make generate
 Below is a list of command-line flags accepted by Katib controller:
 
 | Name                            | Type                      | Default                       | Description                                                                                                            |
-| ------------------------------- | ------------------------- | ----------------------------- | ---------------------------------------------------------------------------------------------------------------------- |
+|---------------------------------|---------------------------|-------------------------------|------------------------------------------------------------------------------------------------------------------------|
 | enable-grpc-probe-in-suggestion | bool                      | true                          | Enable grpc probe in suggestions                                                                                       |
 | experiment-suggestion-name      | string                    | "default"                     | The implementation of suggestion interface in experiment controller                                                    |
 | metrics-addr                    | string                    | ":8080"                       | The address the metric endpoint binds to                                                                               |
@@ -68,6 +68,14 @@ Below is a list of command-line flags accepted by Katib controller:
 | enable-leader-election          | bool                      | false                         | Enable leader election for katib-controller. Enabling this will ensure there is only one active katib-controller.      |
 | leader-election-id              | string                    | "3fbc96e9.katib.kubeflow.org" | The ID for leader election.                                                                                            |
 
+## DB Manager Flags
+
+Below is a list of command-line flags accepted by Katib DB Manager:
+
+| Name            | Type          | Default | Description                                             |
+|-----------------|---------------|---------|---------------------------------------------------------|
+| connect-timeout | time.Duration | 60s     | Timeout before calling error during database connection |
+
 ## Workflow design
 
 Please see [workflow-design.md](./workflow-design.md).

diff --git a/pkg/db/v1beta1/common/const.go b/pkg/db/v1beta1/common/const.go
@@ -20,7 +20,6 @@ import "time"
 
 const (
 	ConnectInterval = 5 * time.Second
-	ConnectTimeout  = 60 * time.Second
 
 	DBUserEnvName     = "DB_USER"
 	DBNameEnvName     = "DB_NAME"

diff --git a/pkg/db/v1beta1/db.go b/pkg/db/v1beta1/db.go
@@ -18,21 +18,22 @@ package db
 
 import (
 	"errors"
+	"time"
 
 	"github.com/kubeflow/katib/pkg/db/v1beta1/common"
 	"github.com/kubeflow/katib/pkg/db/v1beta1/mysql"
 	"github.com/kubeflow/katib/pkg/db/v1beta1/postgres"
 	"k8s.io/klog"
 )
 
-func NewKatibDBInterface(dbName string) (common.KatibDBInterface, error) {
+func NewKatibDBInterface(dbName string, connectTimeout time.Duration) (common.KatibDBInterface, error) {
 
 	if dbName == common.MySqlDBNameEnvValue {
 		klog.Info("Using MySQL")
-		return mysql.NewDBInterface()
+		return mysql.NewDBInterface(connectTimeout)
 	} else if dbName == common.PostgresSQLDBNameEnvValue {
 		klog.Info("Using Postgres")
-		return postgres.NewDBInterface()
+		return postgres.NewDBInterface(connectTimeout)
 	}
 	return nil, errors.New("Invalid DB Name")
 }
diff --git a/pkg/db/v1beta1/mysql/mysql.go b/pkg/db/v1beta1/mysql/mysql.go
@@ -73,8 +73,8 @@ func NewWithSQLConn(db *sql.DB) (common.KatibDBInterface, error) {
 	return d, nil
 }
 
-func NewDBInterface() (common.KatibDBInterface, error) {
-	db, err := common.OpenSQLConn(dbDriver, getDbName(), common.ConnectInterval, common.ConnectTimeout)
+func NewDBInterface(connectTimeout time.Duration) (common.KatibDBInterface, error) {
+	db, err := common.OpenSQLConn(dbDriver, getDbName(), common.ConnectInterval, connectTimeout)
 	if err != nil {
 		return nil, fmt.Errorf("DB open failed: %v", err)
 	}

diff --git a/pkg/db/v1beta1/postgres/postgres.go b/pkg/db/v1beta1/postgres/postgres.go
@@ -49,7 +49,7 @@ func getDbName() string {
 		common.PostgreSQLDBHostEnvName, common.DefaultPostgreSQLHost)
 	dbPort := env.GetEnvOrDefault(
 		common.PostgreSQLDBPortEnvName, common.DefaultPostgreSQLPort)
-	dbName := env.GetEnvOrDefault(common.DefaultPostgreSQLDatabase,
+	dbName := env.GetEnvOrDefault(common.PostgreSQLDatabase,
 		common.DefaultPostgreSQLDatabase)
 
 	psqlInfo := fmt.Sprintf("host=%s port=%s user=%s "+
@@ -59,8 +59,8 @@ func getDbName() string {
 	return psqlInfo
 }
 
-func NewDBInterface() (common.KatibDBInterface, error) {
-	db, err := common.OpenSQLConn(dbDriver, getDbName(), common.ConnectInterval, common.ConnectTimeout)
+func NewDBInterface(connectTimeout time.Duration) (common.KatibDBInterface, error) {
+	db, err := common.OpenSQLConn(dbDriver, getDbName(), common.ConnectInterval, connectTimeout)
 	if err != nil {
 		return nil, fmt.Errorf("DB open failed: %v", err)
 	}