migration ran

Signed-off-by: Yee Hing Tong <[email protected]>

flyteartifacts/cmd/main.go

@@ -14,7 +14,8 @@ func main() {
 	logger.Infof(ctx, "Beginning Flyte Artifacts Service")
 	rootCmd := sharedCmd.NewRootCmd("artifacts", server.GrpcRegistrationHook, server.HttpRegistrationHook)
 	migs := server.GetMigrations(ctx)
-	rootCmd.AddCommand(sharedCmd.NewMigrateCmd(migs))
+	initializationSql := "create extension if not exists hstore;"
+	rootCmd.AddCommand(sharedCmd.NewMigrateCmd(migs, initializationSql))
 	err := rootCmd.ExecuteContext(ctx)
 	if err != nil {
 		panic(err)

flyteartifacts/cmd/shared/migrate.go

@@ -8,12 +8,12 @@ import (
 )
 
 // NewMigrateCmd represents the migrate command
-func NewMigrateCmd(migs []*gormigrate.Migration) *cobra.Command {
+func NewMigrateCmd(migs []*gormigrate.Migration, initializationSql string) *cobra.Command {
 	return &cobra.Command{
 		Use:   "migrate",
 		Short: "This command will run all the migrations for the database",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return database.Migrate(context.Background(), migs)
+			return database.Migrate(context.Background(), migs, initializationSql)
 		},
 	}
 }

flyteartifacts/pkg/server/migrations.go

@@ -21,7 +21,7 @@ var Migrations = []*gormigrate.Migration{
 				gorm.Model
 				ArtifactsKey
 				Version    string          `gorm:"type:varchar(255);index:idx_artifact_version"`
-				Partitions postgres.Hstore `gorm:"index:idx_artifact_partitions"`
+				Partitions postgres.Hstore `gorm:"type:hstore;index:idx_artifact_partitions"`
 
 				LiteralType  []byte `gorm:"not null"`
 				LiteralValue []byte `gorm:"not null"`

flytestdlib/database/db.go

@@ -102,12 +102,17 @@ func withDB(ctx context.Context, do func(db *gorm.DB) error) error {
 }
 
 // Migrate runs all configured migrations
-func Migrate(ctx context.Context, migrations []*gormigrate.Migration) error {
+func Migrate(ctx context.Context, migrations []*gormigrate.Migration, initializationSql string) error {
 	if migrations == nil || len(migrations) == 0 {
 		logger.Infof(ctx, "No migrations to run")
 		return nil
 	}
 	return withDB(ctx, func(db *gorm.DB) error {
+		tx := db.Exec(initializationSql)
+		if tx.Error != nil {
+			return tx.Error
+		}
+
 		m := gormigrate.New(db, gormigrate.DefaultOptions, migrations)
 		if err := m.Migrate(); err != nil {
 			return fmt.Errorf("database migration failed: %v", err)

flytestdlib/database/postgres.go

@@ -9,7 +9,7 @@ import (
 	"strings"
 
 	"github.com/flyteorg/flyte/flytestdlib/logger"
-	"github.com/jackc/pgconn"
+	"github.com/jackc/pgx/v5/pgconn"
 	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 )
@@ -61,17 +61,12 @@ func CreatePostgresDbIfNotExists(ctx context.Context, gormConfig *gorm.Config, p
 	if !isPgErrorWithCode(err, pqInvalidDBCode) {
 		return nil, err
 	}
-	fmt.Println("=============================...........===")
-
 	logger.Warningf(ctx, "Database [%v] does not exist", pgConfig.DbName)
 
 	// Every postgres installation includes a 'postgres' database by default. We connect to that now in order to
 	// initialize the user-specified database.
 	defaultDbPgConfig := pgConfig
 	defaultDbPgConfig.DbName = defaultDB
-	x := getPostgresDsn(ctx, defaultDbPgConfig)
-	fmt.Println(x)
-	fmt.Println("==============================^^===")
 	defaultDBDialector := postgres.Open(getPostgresDsn(ctx, defaultDbPgConfig))
 	gormDb, err = gorm.Open(defaultDBDialector, gormConfig)
 	if err != nil {
@@ -96,9 +91,9 @@ func CreatePostgresDbIfNotExists(ctx context.Context, gormConfig *gorm.Config, p
 }
 
 func isPgErrorWithCode(err error, code string) bool {
+	// Make sure the pgconn you're using is "github.com/jackc/pgx/v5/pgconn"
+	// See https://github.com/go-gorm/gorm/issues/4135
 	pgErr := &pgconn.PgError{}
-	xx := errors.As(err, &pgErr)
-	fmt.Println(xx)
 	if !errors.As(err, &pgErr) {
 		// err chain does not contain a pgconn.PgError
 		return false