Add additional DB connection properties (caraml-dev#86)

* Add additional DB connection properties

* Update unit test cases

* Run gofmt on code

Author: krithika369

api/.gitignore

@@ -0,0 +1 @@
+cover.out

api/config/config.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/caraml-dev/mlp/api/models/v2"
 	"github.com/go-playground/validator/v10"
@@ -55,6 +56,11 @@ type DatabaseConfig struct {
 	Password      string `validate:"required"`
 	Database      string `validate:"required"`
 	MigrationPath string `validate:"required,url"`
+
+	ConnMaxIdleTime time.Duration
+	ConnMaxLifetime time.Duration
+	MaxIdleConns    int
+	MaxOpenConns    int
 }
 
 type AuthorizationConfig struct {

api/config/config_test.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/caraml-dev/mlp/api/config"
 	"github.com/caraml-dev/mlp/api/models/v2"
@@ -34,6 +35,9 @@ func envSetter(envs map[string]string) (closer func()) {
 }
 
 func TestLoad(t *testing.T) {
+	oneSecond, _ := time.ParseDuration("1s")
+	twoSeconds, _ := time.ParseDuration("2s")
+
 	suite := map[string]struct {
 		configs  []string
 		env      map[string]string
@@ -56,11 +60,15 @@ func TestLoad(t *testing.T) {
 					Enabled: false,
 				},
 				Database: &config.DatabaseConfig{
-					Host:          "localhost",
-					Port:          5432,
-					Database:      "mlp",
-					User:          "mlp",
-					MigrationPath: "file://db-migrations",
+					Host:            "localhost",
+					Port:            5432,
+					Database:        "mlp",
+					User:            "mlp",
+					MigrationPath:   "file://db-migrations",
+					ConnMaxIdleTime: oneSecond,
+					ConnMaxLifetime: twoSeconds,
+					MaxIdleConns:    10,
+					MaxOpenConns:    20,
 				},
 				Mlflow: &config.MlflowConfig{},
 				Docs:   []config.Documentation{},
@@ -90,12 +98,16 @@ func TestLoad(t *testing.T) {
 					Enabled: false,
 				},
 				Database: &config.DatabaseConfig{
-					Host:          "localhost",
-					Port:          5432,
-					Database:      "mlp",
-					User:          "mlp",
-					Password:      "secret",
-					MigrationPath: "file://db-migrations",
+					Host:            "localhost",
+					Port:            5432,
+					Database:        "mlp",
+					User:            "mlp",
+					Password:        "secret",
+					MigrationPath:   "file://db-migrations",
+					ConnMaxIdleTime: oneSecond,
+					ConnMaxLifetime: twoSeconds,
+					MaxIdleConns:    10,
+					MaxOpenConns:    20,
 				},
 				Mlflow: &config.MlflowConfig{},
 				Docs:   []config.Documentation{},
@@ -152,12 +164,16 @@ func TestLoad(t *testing.T) {
 					KetoServerURL: "http://localhost:4466",
 				},
 				Database: &config.DatabaseConfig{
-					Host:          "localhost",
-					Port:          5432,
-					Database:      "mlp",
-					User:          "mlp",
-					Password:      "secret",
-					MigrationPath: "file://db-migrations",
+					Host:            "localhost",
+					Port:            5432,
+					Database:        "mlp",
+					User:            "mlp",
+					Password:        "secret",
+					MigrationPath:   "file://db-migrations",
+					ConnMaxIdleTime: oneSecond,
+					ConnMaxLifetime: twoSeconds,
+					MaxIdleConns:    10,
+					MaxOpenConns:    20,
 				},
 				Docs: []config.Documentation{
 					{

api/config/testdata/config-1.yaml

@@ -1,5 +1,9 @@
 database:
   user: mlp
+  connMaxIdleTime: 1s
+  connMaxLifetime: 2s
+  maxIdleConns: 10
+  maxOpenConns: 20
 
 streams:
   stream-1:

api/database/database.go

@@ -31,6 +31,17 @@ func InitDB(dbCfg *config.DatabaseConfig) (*gorm.DB, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	// Get the underlying SQL DB and apply connection properties
+	sqlDB := db.DB()
+	if sqlDB == nil {
+		return nil, errors.New("Failed to get underlying database connection")
+	}
+	sqlDB.SetConnMaxIdleTime(dbCfg.ConnMaxIdleTime)
+	sqlDB.SetConnMaxLifetime(dbCfg.ConnMaxLifetime)
+	sqlDB.SetMaxIdleConns(dbCfg.MaxIdleConns)
+	sqlDB.SetMaxOpenConns(dbCfg.MaxOpenConns)
+
 	db.LogMode(false)
 	err = runDBMigration(db, dbCfg.MigrationPath)
 	if err != nil {