Upgrade gorm to v1.21.9 (flyteorg#42)

Signed-off-by: Daniel Rammer <[email protected]>

cmd/entrypoints/migrate.go

@@ -9,7 +9,6 @@ import (
 
 	"context"
 
-	_ "github.com/jinzhu/gorm/dialects/postgres" // Required to import database driver.
 	"github.com/spf13/cobra"
 )
 
@@ -73,7 +72,10 @@ var migrateCmd = &cobra.Command{
 		logger.Infof(ctx, "Created DB connection.")
 
 		// 	TODO: checkpoints for migrations
-		dbHandle.Migrate()
+		if err := dbHandle.Migrate(ctx); err != nil {
+			logger.Errorf(ctx, "Failed to migrate. err: %v", err)
+			panic(err)
+		}
 		logger.Infof(ctx, "Ran DB migration successfully.")
 	},
 }

datacatalog_config.yaml

@@ -30,3 +30,4 @@ database:
   host: localhost
   dbname: datacatalog
   options: "sslmode=disable"
+  log_level: 5

go.mod

@@ -8,15 +8,11 @@ require (
 	github.com/flyteorg/flytestdlib v0.3.13
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
 	github.com/golang/protobuf v1.4.3
-	github.com/jackc/pgproto3/v2 v2.0.7 // indirect
-	github.com/jinzhu/gorm v1.9.16
 	github.com/lib/pq v1.3.0
 	github.com/mitchellh/mapstructure v1.4.1
 	github.com/spf13/cobra v1.1.1
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.7.0
-	golang.org/x/crypto v0.0.0-20210513164829-c07d793c2f9a // indirect
-	golang.org/x/text v0.3.6 // indirect
 	google.golang.org/grpc v1.36.0
 	gorm.io/driver/postgres v1.1.0
 	gorm.io/gorm v1.21.9

pkg/common/filters.go

@@ -2,8 +2,8 @@ package common
 
 // Common constants and types for Filtering
 const (
-	DefaultPageOffset = uint32(0)
-	MaxPageLimit      = uint32(50)
+	DefaultPageOffset = 0
+	MaxPageLimit      = 50
 )
 
 // Common Entity types that can be used on any filters

pkg/manager/impl/artifact_manager.go

@@ -276,7 +276,7 @@ func (m *artifactManager) ListArtifacts(ctx context.Context, request *datacatalo
 		artifact.Data = artifactDataList
 	}
 
-	token := strconv.Itoa(int(listInput.Offset) + len(artifactsList))
+	token := strconv.Itoa(listInput.Offset + len(artifactsList))
 
 	logger.Debugf(ctx, "Listed %v matching artifacts successfully", len(artifactsList))
 	m.systemMetrics.listSuccessCounter.Inc(ctx)

pkg/manager/impl/dataset_manager.go

@@ -182,7 +182,7 @@ func (dm *datasetManager) ListDatasets(ctx context.Context, request *datacatalog
 		return nil, errors.NewCollectedErrors(codes.Internal, transformerErrs)
 	}
 
-	token := strconv.Itoa(int(listInput.Offset) + len(datasetList))
+	token := strconv.Itoa(listInput.Offset + len(datasetList))
 
 	logger.Debugf(ctx, "Listed %v matching datasets successfully", len(datasetList))
 	dm.systemMetrics.listSuccessCounter.Inc(ctx)

pkg/repositories/config/database.go

@@ -1,5 +1,7 @@
 package config
 
+import "gorm.io/gorm/logger"
+
 //go:generate pflags DbConfigSection
 
 // This struct corresponds to the  database section of in the config
@@ -12,7 +14,8 @@ type DbConfigSection struct {
 	Password     string `json:"password"`
 	PasswordPath string `json:"passwordPath"`
 	// See http://gorm.io/docs/connecting_to_the_database.html for available options passed, in addition to the above.
-	ExtraOptions string `json:"options"`
+	ExtraOptions string          `json:"options"`
+	LogLevel     logger.LogLevel `json:"log_level"`
 }
 
 // Database config. Contains values necessary to open a database connection.

pkg/repositories/config/postgres.go

@@ -3,30 +3,29 @@ package config
 import (
 	"fmt"
 
+	"gorm.io/gorm/logger"
+
 	"github.com/flyteorg/flytestdlib/promutils"
 
-	"github.com/jinzhu/gorm"
-	_ "github.com/jinzhu/gorm/dialects/postgres" // Required to import database driver.
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
 )
 
 const Postgres = "postgres"
 
 // Generic interface for providing a config necessary to open a database connection.
 type DbConnectionConfigProvider interface {
-	// Returns the database type. For instance PostgreSQL or MySQL.
-	GetType() string
-	// Returns arguments specific for the database type necessary to open a database connection.
-	GetArgs() string
-	// Enables verbose logging.
-	WithDebugModeEnabled()
-	// Disables verbose logging.
-	WithDebugModeDisabled()
-	// Returns whether verbose logging is enabled or not.
-	IsDebug() bool
+	// Returns database dialector
+	GetDialector() gorm.Dialector
+
+	GetDBConfig() DbConfig
+
+	GetDSN() string
 }
 
 type BaseConfig struct {
-	IsDebug bool
+	LogLevel                                 logger.LogLevel `json:"log_level"`
+	DisableForeignKeyConstraintWhenMigrating bool
 }
 
 // PostgreSQL implementation for DbConnectionConfigProvider.
@@ -43,11 +42,7 @@ func NewPostgresConfigProvider(config DbConfig, scope promutils.Scope) DbConnect
 	}
 }
 
-func (p *PostgresConfigProvider) GetType() string {
-	return Postgres
-}
-
-func (p *PostgresConfigProvider) GetArgs() string {
+func (p *PostgresConfigProvider) GetDSN() string {
 	if p.config.Password == "" {
 		// Switch for development
 		return fmt.Sprintf("host=%s port=%d dbname=%s user=%s sslmode=disable",
@@ -57,25 +52,23 @@ func (p *PostgresConfigProvider) GetArgs() string {
 		p.config.Host, p.config.Port, p.config.DbName, p.config.User, p.config.Password, p.config.ExtraOptions)
 }
 
-func (p *PostgresConfigProvider) WithDebugModeEnabled() {
-	p.config.IsDebug = true
-}
-
-func (p *PostgresConfigProvider) WithDebugModeDisabled() {
-	p.config.IsDebug = false
+func (p *PostgresConfigProvider) GetDialector() gorm.Dialector {
+	return postgres.Open(p.GetDSN())
 }
 
-func (p *PostgresConfigProvider) IsDebug() bool {
-	return p.config.IsDebug
+func (p *PostgresConfigProvider) GetDBConfig() DbConfig {
+	return p.config
 }
 
 // Opens a connection to the database specified in the config.
 // You must call CloseDbConnection at the end of your session!
 func OpenDbConnection(config DbConnectionConfigProvider) (*gorm.DB, error) {
-	db, err := gorm.Open(config.GetType(), config.GetArgs())
+	db, err := gorm.Open(config.GetDialector(), &gorm.Config{
+		Logger:                                   logger.Default.LogMode(config.GetDBConfig().LogLevel),
+		DisableForeignKeyConstraintWhenMigrating: config.GetDBConfig().DisableForeignKeyConstraintWhenMigrating,
+	})
 	if err != nil {
 		return nil, err
 	}
-	db.LogMode(config.IsDebug())
 	return db, nil
 }

pkg/repositories/config/postgres_test.go

@@ -3,6 +3,8 @@ package config
 import (
 	"testing"
 
+	"gorm.io/gorm/logger"
+
 	mockScope "github.com/flyteorg/flytestdlib/promutils"
 
 	"github.com/stretchr/testify/assert"
@@ -15,9 +17,15 @@ func TestConstructGormArgs(t *testing.T) {
 		DbName:       "postgres",
 		User:         "postgres",
 		ExtraOptions: "sslmode=disable",
+		BaseConfig: BaseConfig{
+			LogLevel:                                 3,
+			DisableForeignKeyConstraintWhenMigrating: true,
+		},
 	}, mockScope.NewTestScope())
 
-	assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres sslmode=disable", postgresConfigProvider.GetArgs())
+	assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres sslmode=disable", postgresConfigProvider.GetDSN())
+	assert.Equal(t, logger.LogLevel(3), postgresConfigProvider.GetDBConfig().LogLevel)
+	assert.Equal(t, true, postgresConfigProvider.GetDBConfig().DisableForeignKeyConstraintWhenMigrating)
 }
 
 func TestConstructGormArgsWithPassword(t *testing.T) {
@@ -30,7 +38,7 @@ func TestConstructGormArgsWithPassword(t *testing.T) {
 		ExtraOptions: "sslmode=enable",
 	}, mockScope.NewTestScope())
 
-	assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres password=pass sslmode=enable", postgresConfigProvider.GetArgs())
+	assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres password=pass sslmode=enable", postgresConfigProvider.GetDSN())
 }
 
 func TestConstructGormArgsWithPasswordNoExtra(t *testing.T) {
@@ -42,5 +50,5 @@ func TestConstructGormArgsWithPasswordNoExtra(t *testing.T) {
 		Password: "pass",
 	}, mockScope.NewTestScope())
 
-	assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres password=pass ", postgresConfigProvider.GetArgs())
+	assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres password=pass ", postgresConfigProvider.GetDSN())
 }

pkg/repositories/errors/postgres.go

@@ -4,9 +4,9 @@ import (
 	"fmt"
 
 	"github.com/flyteorg/datacatalog/pkg/errors"
-	"github.com/jinzhu/gorm"
 	"github.com/lib/pq"
 	"google.golang.org/grpc/codes"
+	"gorm.io/gorm"
 )
 
 // Postgres error codes

pkg/repositories/gormimpl/artifact.go

@@ -3,13 +3,14 @@ package gormimpl
 import (
 	"context"
 
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog"
+
 	"github.com/flyteorg/datacatalog/pkg/common"
 	"github.com/flyteorg/datacatalog/pkg/repositories/errors"
 	"github.com/flyteorg/datacatalog/pkg/repositories/interfaces"
 	"github.com/flyteorg/datacatalog/pkg/repositories/models"
-	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog"
 	"github.com/flyteorg/flytestdlib/promutils"
-	"github.com/jinzhu/gorm"
+	"gorm.io/gorm"
 )
 
 type artifactRepo struct {
@@ -65,19 +66,20 @@ func (h *artifactRepo) Get(ctx context.Context, in models.ArtifactKey) (models.A
 		)
 
 	if result.Error != nil {
+		if result.Error.Error() == gorm.ErrRecordNotFound.Error() {
+			return models.Artifact{}, errors.GetMissingEntityError("Artifact", &datacatalog.Artifact{
+				Dataset: &datacatalog.DatasetID{
+					Project: in.DatasetProject,
+					Domain:  in.DatasetDomain,
+					Name:    in.DatasetName,
+					Version: in.DatasetVersion,
+				},
+				Id: in.ArtifactID,
+			})
+		}
+
 		return models.Artifact{}, h.errorTransformer.ToDataCatalogError(result.Error)
 	}
-	if result.RecordNotFound() {
-		return models.Artifact{}, errors.GetMissingEntityError("Artifact", &datacatalog.Artifact{
-			Dataset: &datacatalog.DatasetID{
-				Project: in.DatasetProject,
-				Domain:  in.DatasetDomain,
-				Name:    in.DatasetName,
-				Version: in.DatasetVersion,
-			},
-			Id: in.ArtifactID,
-		})
-	}
 
 	return artifact, nil
 }