faster migrations (#9060)

optimize migrations on metrics table (landed as part of #9032)
determined-ai · Mar 27, 2024 · a07f0fb · a07f0fb
1 parent e8dba6d
commit a07f0fb
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 55 deletions.
diff --git a/master/static/migrations/20240325144650_generic-metrics-nullable-batches.tx.down.sql b/master/static/migrations/20240325144650_generic-metrics-nullable-batches.tx.down.sql
diff --git a/master/static/migrations/20240325144650_generic-metrics-nullable-batches.tx.up.sql b/master/static/migrations/20240325144650_generic-metrics-nullable-batches.tx.up.sql
diff --git a/master/static/migrations/20240325144710_generic-metrics-text-partition-type.tx.down.sql b/master/static/migrations/20240325144710_generic-metrics-text-partition-type.tx.down.sql
@@ -1,5 +1,6 @@
 /*
- Rollback changing `partition_type` on `metrics` from ENUM type to TEXT.
+ Rollback changing `partition_type` on `metrics` from ENUM type to TEXT,
+ dropping NOT NULL on `total_batches`
 
  Partition keys cannot be modified without dropping and recreating the table.
  */
@@ -52,7 +53,7 @@ CREATE TABLE metrics (
     trial_id integer NOT NULL,
     end_time timestamp with time zone,
     metrics jsonb,
-    total_batches integer,
+    total_batches integer NOT NULL DEFAULT 0,
     trial_run_id integer NOT NULL DEFAULT 0,
     archived boolean NOT NULL DEFAULT false,
     id integer NOT NULL DEFAULT nextval('metrics_id_seq'),

diff --git a/master/static/migrations/20240325144710_generic-metrics-text-partition-type.tx.up.sql b/master/static/migrations/20240325144710_generic-metrics-text-partition-type.tx.up.sql
@@ -1,5 +1,6 @@
 /*
- Change `partition_type` on `metrics` from ENUM type to TEXT.
+  Drop 'NOT NULL' on `metrics.total_batches`.
+  Change `partition_type` on `metrics` from ENUM type to TEXT.
 
  Partition keys cannot be modified without dropping and recreating the table.
  */
@@ -40,7 +41,7 @@ SELECT setval(
     true
 );
 
--- Re-create table with `partition_type` TEXT.
+-- Re-create table with `partition_type` TEXT and nullable `total_batches`
 CREATE TABLE metrics (
     trial_id integer NOT NULL,
     end_time timestamp with time zone,
@@ -55,14 +56,18 @@ CREATE TABLE metrics (
 
 
 -- Modify child partitions to have `partition_type` TEXT and set defaults.
-ALTER TABLE raw_steps ALTER COLUMN partition_type TYPE text;
-ALTER TABLE raw_steps ALTER COLUMN partition_type SET DEFAULT 'TRAINING';
+-- Drop and recreate the columns, since it's faster and we already know the values.
+ALTER TABLE raw_steps DROP COLUMN partition_type;
+ALTER TABLE raw_steps ADD COLUMN partition_type TEXT DEFAULT 'TRAINING';
+ALTER TABLE raw_steps ALTER COLUMN partition_type SET NOT NULL;
 
-ALTER TABLE raw_validations ALTER COLUMN partition_type TYPE text;
-ALTER TABLE raw_validations ALTER COLUMN partition_type SET DEFAULT 'VALIDATION';
+ALTER TABLE raw_validations DROP COLUMN partition_type;
+ALTER TABLE raw_validations ADD COLUMN partition_type TEXT DEFAULT 'VALIDATION';
+ALTER TABLE raw_validations ALTER COLUMN partition_type SET NOT NULL;
 
-ALTER TABLE generic_metrics ALTER COLUMN partition_type TYPE text;
-ALTER TABLE generic_metrics ALTER COLUMN partition_type SET DEFAULT 'GENERIC';
+ALTER TABLE generic_metrics DROP COLUMN partition_type;
+ALTER TABLE generic_metrics ADD COLUMN partition_type TEXT DEFAULT 'GENERIC';
+ALTER TABLE generic_metrics ALTER COLUMN partition_type SET NOT NULL;
 
 
 -- Drop `metric_partition_type` enum.