Read database config from Flytestdlib (flyteorg#380)

* Read database config from datacatalog when failed to register config Signed-off-by: Kevin Su <[email protected]> * Updated dependency Signed-off-by: Kevin Su <[email protected]> * Fixed tests Signed-off-by: Kevin Su <[email protected]> * Fixed tests Signed-off-by: Kevin Su <[email protected]> * Fixed lint Signed-off-by: Kevin Su <[email protected]> * address comment Signed-off-by: Kevin Su <[email protected]> * Fixed tests Signed-off-by: Kevin Su <[email protected]> * Fixed tests error Signed-off-by: Kevin Su <[email protected]>
eapolinario · Apr 4, 2022 · 6930b25 · 6930b25
1 parent b9c0da1
commit 6930b25
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 25 deletions.
diff --git a/go.mod b/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/flyteorg/flyteidl v0.24.10
 	github.com/flyteorg/flyteplugins v0.10.16
 	github.com/flyteorg/flytepropeller v0.16.36
-	github.com/flyteorg/flytestdlib v0.4.16
+	github.com/flyteorg/flytestdlib v0.4.17
 	github.com/flyteorg/stow v0.3.1
 	github.com/ghodss/yaml v1.0.0
 	github.com/go-gormigrate/gormigrate/v2 v2.0.0
@@ -59,15 +59,18 @@ require (
 	sigs.k8s.io/controller-runtime v0.8.3
 )
 
+require (
+	github.com/Azure/azure-sdk-for-go/sdk/azcore v0.21.1 // indirect
+	github.com/Azure/azure-sdk-for-go/sdk/internal v0.8.3 // indirect
+	github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v0.3.0 // indirect
+)
+
 require (
 	cloud.google.com/go v0.100.2 // indirect
 	cloud.google.com/go/compute v0.1.0 // indirect
 	cloud.google.com/go/kms v1.2.0 // indirect
 	cloud.google.com/go/pubsub v1.10.1 // indirect
 	github.com/Azure/azure-sdk-for-go v62.3.0+incompatible // indirect
-	github.com/Azure/azure-sdk-for-go/sdk/azcore v0.21.1 // indirect
-	github.com/Azure/azure-sdk-for-go/sdk/internal v0.8.3 // indirect
-	github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v0.3.0 // indirect
 	github.com/Azure/go-autorest v14.2.0+incompatible // indirect
 	github.com/Azure/go-autorest/autorest v0.11.18 // indirect
 	github.com/Azure/go-autorest/autorest/adal v0.9.13 // indirect
@@ -110,7 +113,7 @@ require (
 	github.com/jackc/pgtype v1.9.0 // indirect
 	github.com/jackc/pgx/v4 v4.14.0 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
-	github.com/jinzhu/now v1.1.3 // indirect
+	github.com/jinzhu/now v1.1.4 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/kelseyhightower/envconfig v1.4.0 // indirect

diff --git a/go.sum b/go.sum
@@ -367,8 +367,8 @@ github.com/flyteorg/flytepropeller v0.16.36 h1:5uE8JsutrPVyLVrRJ8BgvhZUOmTBFkEkn
 github.com/flyteorg/flytepropeller v0.16.36/go.mod h1:DGCjQSRp8VYOBH56aQyAZfNf1Vgh+GNpwQL7uhottYM=
 github.com/flyteorg/flytestdlib v0.3.13/go.mod h1:Tz8JCECAbX6VWGwFT6cmEQ+RJpZ/6L9pswu3fzWs220=
 github.com/flyteorg/flytestdlib v0.4.13/go.mod h1:fv1ar34LJLMTaf0tbfetisLykUlARi7rP+NQTUn6QQs=
-github.com/flyteorg/flytestdlib v0.4.16 h1:r4dCPUOqoE9xCAhOw9KDB7O6cBoCxyEtepIWYcj93H0=
-github.com/flyteorg/flytestdlib v0.4.16/go.mod h1:WA5Y4hrcgD0ybGOKJVOQ4sP8q7NLRV+S5SWOlH0axgM=
+github.com/flyteorg/flytestdlib v0.4.17 h1:C/Q0vcQNlEdX1p0TZyOZOqjSB30Sd5gqBQsuMGUQqmA=
+github.com/flyteorg/flytestdlib v0.4.17/go.mod h1:GsZJqXrtgSIb5D3i75+Ud+1ufdPVEwvQegNdo22DPMo=
 github.com/flyteorg/stow v0.3.1 h1:cBMbWl03Gsy5KoA5mutUYTuYpqtT7Pb8+ANGCLnmFEs=
 github.com/flyteorg/stow v0.3.1/go.mod h1:HBld7ud0i4khMHwJjkO8v+NSP7ddKa/ruhf4I8fliaA=
 github.com/fogleman/gg v1.2.1-0.20190220221249-0403632d5b90/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k=
@@ -995,8 +995,9 @@ github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD
 github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
 github.com/jinzhu/now v1.1.1/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
 github.com/jinzhu/now v1.1.2/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
-github.com/jinzhu/now v1.1.3 h1:PlHq1bSCSZL9K0wUhbm2pGLoTWs2GwVhsP6emvGV/ZI=
 github.com/jinzhu/now v1.1.3/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
+github.com/jinzhu/now v1.1.4 h1:tHnRBy1i5F2Dh8BAFxqFzxKqqvezXrL2OW1TnX+Mlas=
+github.com/jinzhu/now v1.1.4/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
 github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af/go.mod h1:Nht3zPeWKUH0NzdCt2Blrr5ys8VGpn0CEB0cQHVjt7k=
 github.com/jmespath/go-jmespath v0.3.0/go.mod h1:9QtRXoHjLGCJ5IBSaohpXITPlowMeeYCZ7fLUTSywik=
 github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=

diff --git a/pkg/runtime/application_config_provider.go b/pkg/runtime/application_config_provider.go
@@ -1,36 +1,25 @@
 package runtime
 
 import (
-	"time"
+	"fmt"
 
 	"github.com/flyteorg/flyteadmin/pkg/common"
 	"github.com/flyteorg/flyteadmin/pkg/runtime/interfaces"
 	"github.com/flyteorg/flytestdlib/config"
+	"github.com/flyteorg/flytestdlib/database"
 )
 
-const database = "database"
 const flyteAdmin = "flyteadmin"
 const scheduler = "scheduler"
 const remoteData = "remoteData"
 const notifications = "notifications"
 const domains = "domains"
 const externalEvents = "externalEvents"
 const metricPort = 10254
-const postgres = "postgres"
 
 const KB = 1024
 const MB = KB * KB
 
-var databaseConfig = config.MustRegisterSection(database, &interfaces.DbConfig{
-	DeprecatedPort:         5432,
-	DeprecatedUser:         postgres,
-	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,
 	MetricsScope:          "flyte:",
@@ -90,7 +79,17 @@ var externalEventsConfig = config.MustRegisterSection(externalEvents, &interface
 type ApplicationConfigurationProvider struct{}
 
 func (p *ApplicationConfigurationProvider) GetDbConfig() *interfaces.DbConfig {
-	return databaseConfig.GetConfig().(*interfaces.DbConfig)
+	databaseConfig := database.GetConfig()
+	switch {
+	case !databaseConfig.SQLite.IsEmpty():
+		sqliteConfig := interfaces.SQLiteConfig(databaseConfig.SQLite)
+		return &interfaces.DbConfig{SQLiteConfig: &sqliteConfig}
+	case !databaseConfig.Postgres.IsEmpty():
+		postgresConfig := interfaces.PostgresConfig(databaseConfig.Postgres)
+		return &interfaces.DbConfig{PostgresConfig: &postgresConfig}
+	default:
+		panic(fmt.Errorf("database config cannot be empty"))
+	}
 }
 
 func (p *ApplicationConfigurationProvider) GetTopLevelConfig() *interfaces.ApplicationConfig {

diff --git a/pkg/runtime/config_provider_test.go b/pkg/runtime/config_provider_test.go
@@ -12,21 +12,21 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func initTestConfig() error {
+func initConfig(cfg string) error {
 	pwd, err := os.Getwd()
 	if err != nil {
 		return err
 	}
 
 	configAccessor := viper.NewAccessor(config.Options{
-		SearchPaths: []string{filepath.Join(pwd, "testdata/clusters_config.yaml")},
+		SearchPaths: []string{filepath.Join(pwd, cfg)},
 		StrictMode:  false,
 	})
 	return configAccessor.UpdateConfig(context.TODO())
 }
 
 func TestClusterConfig(t *testing.T) {
-	err := initTestConfig()
+	err := initConfig("testdata/clusters_config.yaml")
 	assert.NoError(t, err)
 
 	configProvider := NewConfigurationProvider()
@@ -49,3 +49,25 @@ func TestClusterConfig(t *testing.T) {
 
 	assert.Equal(t, "file_path", clusters[1].Auth.Type)
 }
+
+func TestPostgresConfig(t *testing.T) {
+	err := initConfig("testdata/postgres_config.yaml")
+	assert.NoError(t, err)
+
+	configProvider := NewConfigurationProvider()
+	dbConfig := configProvider.ApplicationConfiguration().GetDbConfig()
+	assert.Equal(t, 5432, dbConfig.PostgresConfig.Port)
+	assert.Equal(t, "postgres", dbConfig.PostgresConfig.Host)
+	assert.Equal(t, "postgres", dbConfig.PostgresConfig.User)
+	assert.Equal(t, "postgres", dbConfig.PostgresConfig.DbName)
+	assert.Equal(t, "sslmode=disable", dbConfig.PostgresConfig.ExtraOptions)
+}
+
+func TestSqliteConfig(t *testing.T) {
+	err := initConfig("testdata/sqlite_config.yaml")
+	assert.NoError(t, err)
+
+	configProvider := NewConfigurationProvider()
+	dbConfig := configProvider.ApplicationConfiguration().GetDbConfig()
+	assert.Equal(t, "admin.db", dbConfig.SQLiteConfig.File)
+}
diff --git a/pkg/runtime/testdata/postgres_config.yaml b/pkg/runtime/testdata/postgres_config.yaml
@@ -0,0 +1,7 @@
+database:
+  postgres:
+    port: 5432
+    username: postgres
+    host: postgres
+    dbname: postgres
+    options: "sslmode=disable"
diff --git a/pkg/runtime/testdata/sqlite_config.yaml b/pkg/runtime/testdata/sqlite_config.yaml
@@ -0,0 +1,3 @@
+database:
+  sqlite:
+    file: admin.db