From 1199b480cae4fcd32b1687cef6c55476504c2c95 Mon Sep 17 00:00:00 2001
From: Katrina Rogan <katroganGH@gmail.com>
Date: Fri, 8 Apr 2022 12:58:40 -0700
Subject: [PATCH] Configure IgnoreRecordNotFoundError in gorm logger (#396)

Signed-off-by: Katrina Rogan <katroganGH@gmail.com>
---
 pkg/repositories/database.go      | 33 ++++++++++++++++++++++---------
 pkg/repositories/database_test.go | 26 ------------------------
 2 files changed, 24 insertions(+), 35 deletions(-)

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")