Create indices on recently added fields (#16)

G-Research · May 2, 2023 · 4ad03ec · 4ad03ec
1 parent 5441d97
commit 4ad03ec
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/pkg/database/database.go b/pkg/database/database.go
@@ -149,7 +149,7 @@ func ConnectDB(dsn string, slowThreshold time.Duration, poolMax int, reset bool,
 		tx.First(&schemaVersion)
 	}
 
-	if alembicVersion.Version != "97727af70f4d" || schemaVersion.Version != "8073e7e037e5" {
+	if alembicVersion.Version != "97727af70f4d" || schemaVersion.Version != "ed364de02645" {
 		if !migrate && alembicVersion.Version != "" {
 			return fmt.Errorf("unsupported database schema versions alembic %s, fasttrack %s", alembicVersion.Version, schemaVersion.Version)
 		}
@@ -311,6 +311,23 @@ func ConnectDB(dsn string, slowThreshold time.Duration, poolMax int, reset bool,
 					return fmt.Errorf("error migrating database to fasttrack schema 8073e7e037e5: %w", err)
 				}
 
+			case "8073e7e037e5":
+				log.Info("Migrating database to fasttrack schema ed364de02645")
+				if err := DB.Transaction(func(tx *gorm.DB) error {
+					if err := tx.Migrator().CreateIndex(&Run{}, "RowNum"); err != nil {
+						return err
+					}
+					if err := tx.Migrator().CreateIndex(&Metric{}, "Iter"); err != nil {
+						return err
+					}
+					return tx.Model(&SchemaVersion{}).
+						Where("1 = 1").
+						Update("Version", "ed364de02645").
+						Error
+				}); err != nil {
+					return fmt.Errorf("error migrating database to fasttrack schema ed364de02645: %w", err)
+				}
+
 			default:
 				return fmt.Errorf("unsupported database fasttrack schema version %s", schemaVersion.Version)
 			}
@@ -339,7 +356,7 @@ func ConnectDB(dsn string, slowThreshold time.Duration, poolMax int, reset bool,
 				Version: "97727af70f4d",
 			})
 			tx.Create(&SchemaVersion{
-				Version: "8073e7e037e5",
+				Version: "ed364de02645",
 			})
 			tx.Commit()
 			if tx.Error != nil {

diff --git a/pkg/database/model.go b/pkg/database/model.go
@@ -63,7 +63,7 @@ type Run struct {
 	ExperimentID   int32
 	Experiment     Experiment
 	DeletedTime    sql.NullInt64 `gorm:"type:bigint"`
-	RowNum         RowNum
+	RowNum         RowNum        `gorm:"index"`
 	Params         []Param
 	Tags           []Tag
 	Metrics        []Metric
@@ -110,7 +110,7 @@ type Metric struct {
 	RunID     string  `gorm:"column:run_uuid;not null;primaryKey;index"`
 	Step      int64   `gorm:"default:0;not null;primaryKey"`
 	IsNan     bool    `gorm:"default:false;not null;primaryKey"`
-	Iter      int64
+	Iter      int64   `gorm:"index"`
 }
 
 type LatestMetric struct {