1 files changed, 56 insertions, 0 deletions
diff --git a/db/post_migrate/20221220131020_bump_default_partition_id_value_for_ci_tables.rb b/db/post_migrate/20221220131020_bump_default_partition_id_value_for_ci_tables.rb
new file mode 100644
index 00000000000..3d06f02a6d6
--- /dev/null
+++ b/db/post_migrate/20221220131020_bump_default_partition_id_value_for_ci_tables.rb
@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+class BumpDefaultPartitionIdValueForCiTables < Gitlab::Database::Migration[2.1]
+  disable_ddl_transaction!
+
+  TABLES = {
+    ci_build_needs: [:partition_id],
+    ci_build_pending_states: [:partition_id],
+    ci_build_report_results: [:partition_id],
+    ci_build_trace_chunks: [:partition_id],
+    ci_build_trace_metadata: [:partition_id],
+    ci_builds: [:partition_id],
+    ci_builds_runner_session: [:partition_id],
+    ci_job_artifacts: [:partition_id],
+    ci_job_variables: [:partition_id],
+    ci_pending_builds: [:partition_id],
+    ci_pipeline_variables: [:partition_id],
+    ci_pipelines: [:partition_id],
+    ci_running_builds: [:partition_id],
+    ci_sources_pipelines: [:partition_id, :source_partition_id],
+    ci_stages: [:partition_id],
+    ci_unit_test_failures: [:partition_id],
+    p_ci_builds_metadata: [:partition_id]
+  }
+
+  def up
+    change_partitions_default_value(from: 100, to: 101)
+  end
+
+  def down
+    change_partitions_default_value(from: 101, to: 100)
+  end
+
+  private
+
+  def change_partitions_default_value(from:, to:)
+    return unless Gitlab.com?
+
+    TABLES.each do |table_name, columns|
+      next if columns.all? { |column_name| default_value_for(table_name, column_name) == to }
+
+      with_lock_retries do
+        columns.each do |column_name| # rubocop:disable Migration/WithLockRetriesDisallowedMethod
+          change_column_default(table_name, column_name, from: from, to: to)
+        end
+      end
+    end
+  end
+
+  def default_value_for(table_name, column_name)
+    connection
+      .columns(table_name)
+      .find { |column| column.name == column_name.to_s }
+      .default&.to_i
+  end
+end