Skip to content

Commit

Permalink
Add additional DB connection properties (#86)
Browse files Browse the repository at this point in the history
* Add additional DB connection properties

* Update unit test cases

* Run gofmt on code
  • Loading branch information
krithika369 committed May 24, 2023
1 parent 779530a commit f6948ba
Show file tree
Hide file tree
Showing 5 changed files with 55 additions and 17 deletions.
1 change: 1 addition & 0 deletions api/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
cover.out
6 changes: 6 additions & 0 deletions api/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"encoding/json"
"fmt"
"strings"
"time"

"github.com/caraml-dev/mlp/api/models/v2"
"github.com/go-playground/validator/v10"
Expand Down Expand Up @@ -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 {
Expand Down
50 changes: 33 additions & 17 deletions api/config/config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"errors"
"os"
"testing"
"time"

"github.com/caraml-dev/mlp/api/config"
"github.com/caraml-dev/mlp/api/models/v2"
Expand Down Expand Up @@ -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
Expand All @@ -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{},
Expand Down Expand Up @@ -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{},
Expand Down Expand Up @@ -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{
{
Expand Down
4 changes: 4 additions & 0 deletions api/config/testdata/config-1.yaml
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
database:
user: mlp
connMaxIdleTime: 1s
connMaxLifetime: 2s
maxIdleConns: 10
maxOpenConns: 20

streams:
stream-1:
Expand Down
11 changes: 11 additions & 0 deletions api/database/database.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit f6948ba

Please sign in to comment.