Added connection pooling to DB config and added defaults for the same (…

…flyteorg#358) * Added connection pooling to DB config and added defaults for the same Signed-off-by: Prafulla Mahindrakar <[email protected]> * go mod changes for sql lite used for unit tests Signed-off-by: Prafulla Mahindrakar <[email protected]> * linter fixes Signed-off-by: Prafulla Mahindrakar <[email protected]> * test fix Signed-off-by: Prafulla Mahindrakar <[email protected]> * feedback Signed-off-by: Prafulla Mahindrakar <[email protected]> * updated the description for the flags Signed-off-by: Prafulla Mahindrakar <[email protected]>
eapolinario · Mar 2, 2022 · a5014ae · a5014ae
1 parent 7955717
commit a5014ae
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 5 deletions.
diff --git a/flyteadmin/go.mod b/flyteadmin/go.mod
@@ -51,6 +51,7 @@ require (
 	google.golang.org/grpc v1.43.0
 	google.golang.org/protobuf v1.27.1
 	gorm.io/driver/postgres v1.2.3
+	gorm.io/driver/sqlite v1.1.1
 	gorm.io/gorm v1.22.4
 	k8s.io/api v0.20.4
 	k8s.io/apimachinery v0.20.4
@@ -117,6 +118,7 @@ require (
 	github.com/lib/pq v1.10.4 // indirect
 	github.com/mattn/go-colorable v0.1.12 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
+	github.com/mattn/go-sqlite3 v2.0.3+incompatible // indirect
 	github.com/mattn/goveralls v0.0.6 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect

diff --git a/flyteadmin/pkg/repositories/database.go b/flyteadmin/pkg/repositories/database.go
@@ -98,9 +98,26 @@ func GetDB(ctx context.Context, dbConfig *runtimeInterfaces.DbConfig, logConfig
 	default:
 		panic(fmt.Sprintf("Unrecognized database config %v", dbConfig))
 	}
-
-	return gorm.Open(dialector, &gorm.Config{
+	gormDb, err := gorm.Open(dialector, &gorm.Config{
 		Logger:                                   gormLogger.Default.LogMode(logLevel),
 		DisableForeignKeyConstraintWhenMigrating: !dbConfig.EnableForeignKeyConstraintWhenMigrating,
 	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	// Setup connection pool settings
+	return gormDb, setupDbConnectionPool(gormDb, dbConfig)
+}
+
+func setupDbConnectionPool(gormDb *gorm.DB, dbConfig *runtimeInterfaces.DbConfig) error {
+	genericDb, err := gormDb.DB()
+	if err != nil {
+		return err
+	}
+	genericDb.SetConnMaxLifetime(dbConfig.ConnMaxLifeTime.Duration)
+	genericDb.SetMaxIdleConns(dbConfig.MaxIdleConnections)
+	genericDb.SetMaxOpenConns(dbConfig.MaxOpenConnections)
+	return nil
 }
diff --git a/flyteadmin/pkg/repositories/database_test.go b/flyteadmin/pkg/repositories/database_test.go
@@ -3,13 +3,19 @@ package repositories
 import (
 	"context"
 	"io/ioutil"
+	"os"
+	"path/filepath"
 	"testing"
+	"time"
 
 	runtimeInterfaces "github.com/flyteorg/flyteadmin/pkg/runtime/interfaces"
+	"github.com/flyteorg/flytestdlib/config"
 	"github.com/flyteorg/flytestdlib/logger"
-	gormLogger "gorm.io/gorm/logger"
 
 	"github.com/stretchr/testify/assert"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	gormLogger "gorm.io/gorm/logger"
 )
 
 func TestGetGormLogLevel(t *testing.T) {
@@ -90,3 +96,50 @@ func TestGetPostgresDsn(t *testing.T) {
 		assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres password=123abc ", dsn)
 	})
 }
+
+func TestSetupDbConnectionPool(t *testing.T) {
+	t.Run("successful", func(t *testing.T) {
+		gormDb, err := gorm.Open(sqlite.Open(filepath.Join(os.TempDir(), "gorm.db")), &gorm.Config{})
+		assert.Nil(t, err)
+		dbConfig := &runtimeInterfaces.DbConfig{
+			DeprecatedPort:     5432,
+			MaxIdleConnections: 10,
+			MaxOpenConnections: 1000,
+			ConnMaxLifeTime:    config.Duration{Duration: time.Hour},
+		}
+		err = setupDbConnectionPool(gormDb, dbConfig)
+		assert.Nil(t, err)
+		genericDb, err := gormDb.DB()
+		assert.Nil(t, err)
+		assert.Equal(t, genericDb.Stats().MaxOpenConnections, 1000)
+	})
+	t.Run("unset duration", func(t *testing.T) {
+		gormDb, err := gorm.Open(sqlite.Open(filepath.Join(os.TempDir(), "gorm.db")), &gorm.Config{})
+		assert.Nil(t, err)
+		dbConfig := &runtimeInterfaces.DbConfig{
+			DeprecatedPort:     5432,
+			MaxIdleConnections: 10,
+			MaxOpenConnections: 1000,
+		}
+		err = setupDbConnectionPool(gormDb, dbConfig)
+		assert.Nil(t, err)
+		genericDb, err := gormDb.DB()
+		assert.Nil(t, err)
+		assert.Equal(t, genericDb.Stats().MaxOpenConnections, 1000)
+	})
+	t.Run("failed to get DB", func(t *testing.T) {
+		gormDb := &gorm.DB{
+			Config: &gorm.Config{
+				ConnPool: &gorm.PreparedStmtDB{},
+			},
+		}
+		dbConfig := &runtimeInterfaces.DbConfig{
+			DeprecatedPort:     5432,
+			MaxIdleConnections: 10,
+			MaxOpenConnections: 1000,
+			ConnMaxLifeTime:    config.Duration{Duration: time.Hour},
+		}
+		err := setupDbConnectionPool(gormDb, dbConfig)
+		assert.NotNil(t, err)
+	})
+}
diff --git a/flyteadmin/pkg/runtime/application_config_provider.go b/flyteadmin/pkg/runtime/application_config_provider.go
@@ -1,6 +1,8 @@
 package runtime
 
 import (
+	"time"
+
 	"github.com/flyteorg/flyteadmin/pkg/common"
 	"github.com/flyteorg/flyteadmin/pkg/runtime/interfaces"
 	"github.com/flyteorg/flytestdlib/config"
@@ -25,6 +27,9 @@ var databaseConfig = config.MustRegisterSection(database, &interfaces.DbConfig{
 	DeprecatedHost:         postgres,
 	DeprecatedDbName:       postgres,
 	DeprecatedExtraOptions: "sslmode=disable",
+	MaxIdleConnections:     10,
+	MaxOpenConnections:     1000,
+	ConnMaxLifeTime:        config.Duration{Duration: time.Hour},
 })
 var flyteAdminConfig = config.MustRegisterSection(flyteAdmin, &interfaces.ApplicationConfig{
 	ProfilerPort:          metricPort,

diff --git a/flyteadmin/pkg/runtime/interfaces/application_configuration.go b/flyteadmin/pkg/runtime/interfaces/application_configuration.go
@@ -2,6 +2,7 @@ package interfaces
 
 import (
 	"github.com/flyteorg/flytestdlib/config"
+
 	"golang.org/x/time/rate"
 )
 
@@ -17,8 +18,11 @@ type DbConfig struct {
 	DeprecatedExtraOptions string `json:"options" pflag:",deprecated"`
 	DeprecatedDebug        bool   `json:"debug" pflag:",deprecated"`
 
-	EnableForeignKeyConstraintWhenMigrating bool           `json:"enableForeignKeyConstraintWhenMigrating" pflag:",Whether to enable gorm foreign keys when migrating the db"`
-	PostgresConfig                          PostgresConfig `json:"postgres"`
+	EnableForeignKeyConstraintWhenMigrating bool            `json:"enableForeignKeyConstraintWhenMigrating" pflag:",Whether to enable gorm foreign keys when migrating the db"`
+	MaxIdleConnections                      int             `json:"maxIdleConnections" pflag:",maxIdleConnections sets the maximum number of connections in the idle connection pool."`
+	MaxOpenConnections                      int             `json:"maxOpenConnections" pflag:",maxOpenConnections sets the maximum number of open connections to the database."`
+	ConnMaxLifeTime                         config.Duration `json:"connMaxLifeTime" pflag:",sets the maximum amount of time a connection may be reused"`
+	PostgresConfig                          PostgresConfig  `json:"postgres"`
 }
 
 // PostgresConfig includes specific config options for opening a connection to a postgres database.