From 445cd69bac0d292a2649640161ea21e14aae7002 Mon Sep 17 00:00:00 2001
From: Maxime Orefice <morefice@gitlab.com>
Date: Wed, 8 Feb 2023 14:42:53 +0100
Subject: [PATCH 1/2] Fix partition_id for ci_build_needs

Changelog: fixed
---
 ...510_fix_partition_ids_for_ci_build_need.rb | 23 +++++++++
 db/schema_migrations/20230208133510           |  1 +
 ...ix_partition_ids_for_ci_build_need_spec.rb | 51 +++++++++++++++++++
 3 files changed, 75 insertions(+)
 create mode 100644 db/post_migrate/20230208133510_fix_partition_ids_for_ci_build_need.rb
 create mode 100644 db/schema_migrations/20230208133510
 create mode 100644 spec/migrations/20230208133510_fix_partition_ids_for_ci_build_need_spec.rb

diff --git a/db/post_migrate/20230208133510_fix_partition_ids_for_ci_build_need.rb b/db/post_migrate/20230208133510_fix_partition_ids_for_ci_build_need.rb
new file mode 100644
index 00000000000000..a64c34e9cfadb2
--- /dev/null
+++ b/db/post_migrate/20230208133510_fix_partition_ids_for_ci_build_need.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+class FixPartitionIdsForCiBuildNeed < Gitlab::Database::Migration[2.1]
+  disable_ddl_transaction!
+  restrict_gitlab_migration gitlab_schema: :gitlab_ci
+
+  BATCH_SIZE = 100
+
+  def up
+    return unless Gitlab.com?
+
+    define_batchable_model(:ci_build_needs)
+      .where(partition_id: 101)
+      .each_batch(of: BATCH_SIZE) do |batch|
+        batch.update_all(partition_id: 100)
+        sleep rand
+      end
+  end
+
+  def down
+    # no-op
+  end
+end
diff --git a/db/schema_migrations/20230208133510 b/db/schema_migrations/20230208133510
new file mode 100644
index 00000000000000..7e35ba65158eea
--- /dev/null
+++ b/db/schema_migrations/20230208133510
@@ -0,0 +1 @@
+26c4e7a045f8bfe9ceae70038e62838a7015558c219f4112f97863f0c851f7ad
\ No newline at end of file
diff --git a/spec/migrations/20230208133510_fix_partition_ids_for_ci_build_need_spec.rb b/spec/migrations/20230208133510_fix_partition_ids_for_ci_build_need_spec.rb
new file mode 100644
index 00000000000000..75dcce80eaa525
--- /dev/null
+++ b/spec/migrations/20230208133510_fix_partition_ids_for_ci_build_need_spec.rb
@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require_migration!
+
+RSpec.describe FixPartitionIdsForCiBuildNeed, migration: :gitlab_ci, feature_category: :continuous_integration do
+  let(:build_needs) { table(:ci_build_needs, database: :ci) }
+  let(:ci_builds) { table(:ci_builds, database: :ci) }
+  let(:connection) { build_needs.connection }
+
+  around do |example|
+    connection.execute "ALTER TABLE #{build_needs.quoted_table_name} DISABLE TRIGGER ALL;"
+
+    example.run
+
+    connection.execute "ALTER TABLE #{build_needs.quoted_table_name} ENABLE TRIGGER ALL;"
+  end
+
+  before do
+    build = ci_builds.create!(partition_id: 100)
+
+    build_needs.insert_all!([
+      { build_id: build.id, partition_id: 100, name: 'name-100' },
+      { build_id: build.id, partition_id: 101, name: 'name-101' }
+    ])
+  end
+
+  describe '#up' do
+    context 'when on saas' do
+      before do
+        allow(Gitlab).to receive(:com?).and_return(true)
+      end
+
+      it 'fixes partition_id' do
+        expect { migrate! }.not_to raise_error
+
+        expect(build_needs.where(partition_id: 100).count).to eq(2)
+        expect(build_needs.where(partition_id: 101).count).to eq(0)
+      end
+    end
+
+    context 'when on self managed' do
+      it 'does not change partition_id' do
+        expect { migrate! }.not_to raise_error
+
+        expect(build_needs.where(partition_id: 100).count).to eq(1)
+        expect(build_needs.where(partition_id: 101).count).to eq(1)
+      end
+    end
+  end
+end
-- 
GitLab


From 7c69178a5778eeeddd5e566eb2eae4ad94d66dc8 Mon Sep 17 00:00:00 2001
From: Maxime Orefice <morefice@gitlab.com>
Date: Mon, 13 Feb 2023 13:59:57 +0100
Subject: [PATCH 2/2] Increase batch size

---
 .../20230208133510_fix_partition_ids_for_ci_build_need.rb     | 4 ++--
 ...20230208133510_fix_partition_ids_for_ci_build_need_spec.rb | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/db/post_migrate/20230208133510_fix_partition_ids_for_ci_build_need.rb b/db/post_migrate/20230208133510_fix_partition_ids_for_ci_build_need.rb
index a64c34e9cfadb2..d1d1f5cca4ee06 100644
--- a/db/post_migrate/20230208133510_fix_partition_ids_for_ci_build_need.rb
+++ b/db/post_migrate/20230208133510_fix_partition_ids_for_ci_build_need.rb
@@ -4,7 +4,7 @@ class FixPartitionIdsForCiBuildNeed < Gitlab::Database::Migration[2.1]
   disable_ddl_transaction!
   restrict_gitlab_migration gitlab_schema: :gitlab_ci
 
-  BATCH_SIZE = 100
+  BATCH_SIZE = 300
 
   def up
     return unless Gitlab.com?
@@ -13,7 +13,7 @@ def up
       .where(partition_id: 101)
       .each_batch(of: BATCH_SIZE) do |batch|
         batch.update_all(partition_id: 100)
-        sleep rand
+        sleep 0.1
       end
   end
 
diff --git a/spec/migrations/20230208133510_fix_partition_ids_for_ci_build_need_spec.rb b/spec/migrations/20230208133510_fix_partition_ids_for_ci_build_need_spec.rb
index 75dcce80eaa525..c9d1c97123e935 100644
--- a/spec/migrations/20230208133510_fix_partition_ids_for_ci_build_need_spec.rb
+++ b/spec/migrations/20230208133510_fix_partition_ids_for_ci_build_need_spec.rb
@@ -13,6 +13,7 @@
 
     example.run
 
+  ensure
     connection.execute "ALTER TABLE #{build_needs.quoted_table_name} ENABLE TRIGGER ALL;"
   end
 
-- 
GitLab