diff --git a/pkg/repositories/database.go b/pkg/repositories/database.go index 74f72a2bb2..679c3ce09d 100644 --- a/pkg/repositories/database.go +++ b/pkg/repositories/database.go @@ -4,9 +4,11 @@ import ( "context" "fmt" "io/ioutil" + "log" "os" "reflect" "strings" + "time" repoErrors "github.com/flyteorg/flyteadmin/pkg/repositories/errors" "gorm.io/driver/sqlite" @@ -23,27 +25,40 @@ const pqInvalidDBCode = "3D000" const defaultDB = "postgres" // getGormLogLevel converts between the flytestdlib configured log level to the equivalent gorm log level. -func getGormLogLevel(ctx context.Context, logConfig *logger.Config) gormLogger.LogLevel { - if logConfig == nil { +func getGormLogger(ctx context.Context, logConfig *logger.Config) gormLogger.Interface { + logConfigLevel := logger.ErrorLevel + if logConfig != nil { + logConfigLevel = logConfig.Level + } else { logger.Debugf(ctx, "No log config block found, setting gorm db log level to: error") - return gormLogger.Error } - switch logConfig.Level { + var logLevel gormLogger.LogLevel + ignoreRecordNotFoundError := true + switch logConfigLevel { case logger.PanicLevel: fallthrough case logger.FatalLevel: fallthrough case logger.ErrorLevel: - return gormLogger.Error + logLevel = gormLogger.Error case logger.WarnLevel: - return gormLogger.Warn + logLevel = gormLogger.Warn case logger.InfoLevel: fallthrough case logger.DebugLevel: - return gormLogger.Info + logLevel = gormLogger.Info + ignoreRecordNotFoundError = false default: - return gormLogger.Silent + logLevel = gormLogger.Silent } + // Copied from gormLogger.Default initialization. The gormLogger interface only allows modifying the LogLevel + // and not IgnoreRecordNotFoundError. + return gormLogger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), gormLogger.Config{ + SlowThreshold: 200 * time.Millisecond, + LogLevel: logLevel, + IgnoreRecordNotFoundError: ignoreRecordNotFoundError, + Colorful: true, + }) } // Resolves a password value from either a user-provided inline value or a filepath whose contents contain a password. @@ -86,7 +101,7 @@ func GetDB(ctx context.Context, dbConfig *runtimeInterfaces.DbConfig, logConfig panic("Cannot initialize database repository from empty db config") } gormConfig := &gorm.Config{ - Logger: gormLogger.Default.LogMode(getGormLogLevel(ctx, logConfig)), + Logger: getGormLogger(ctx, logConfig), DisableForeignKeyConstraintWhenMigrating: !dbConfig.EnableForeignKeyConstraintWhenMigrating, } diff --git a/pkg/repositories/database_test.go b/pkg/repositories/database_test.go index af81088089..80c5cdae81 100644 --- a/pkg/repositories/database_test.go +++ b/pkg/repositories/database_test.go @@ -16,34 +16,8 @@ import ( "github.com/stretchr/testify/assert" "gorm.io/driver/sqlite" "gorm.io/gorm" - gormLogger "gorm.io/gorm/logger" ) -func TestGetGormLogLevel(t *testing.T) { - assert.Equal(t, gormLogger.Error, getGormLogLevel(context.TODO(), &logger.Config{ - Level: logger.PanicLevel, - })) - assert.Equal(t, gormLogger.Error, getGormLogLevel(context.TODO(), &logger.Config{ - Level: logger.FatalLevel, - })) - assert.Equal(t, gormLogger.Error, getGormLogLevel(context.TODO(), &logger.Config{ - Level: logger.ErrorLevel, - })) - - assert.Equal(t, gormLogger.Warn, getGormLogLevel(context.TODO(), &logger.Config{ - Level: logger.WarnLevel, - })) - - assert.Equal(t, gormLogger.Info, getGormLogLevel(context.TODO(), &logger.Config{ - Level: logger.InfoLevel, - })) - assert.Equal(t, gormLogger.Info, getGormLogLevel(context.TODO(), &logger.Config{ - Level: logger.DebugLevel, - })) - - assert.Equal(t, gormLogger.Error, getGormLogLevel(context.TODO(), nil)) -} - func TestResolvePassword(t *testing.T) { password := "123abc" tmpFile, err := ioutil.TempFile("", "prefix")