Update DynamicModelHelpers to require a connection

Update the methods in DynamicModelHelpers to require a connection that
dynamic models will use to talk to the appropriate database. This change
will ensure that usages of dynamic models are compatible with decomposed
databases.

db/post_migrate/20210415155043_move_container_registry_enabled_to_project_features3.rb

@@ -2,7 +2,6 @@
 
 class MoveContainerRegistryEnabledToProjectFeatures3 < ActiveRecord::Migration[6.0]
   include Gitlab::Database::MigrationHelpers
-  include Gitlab::Database::DynamicModelHelpers
 
   BATCH_SIZE = 21_000
   MIGRATION = 'MoveContainerRegistryEnabledToProjectFeature'

db/post_migrate/20210525075724_clean_up_pending_builds_table.rb

@@ -10,7 +10,7 @@ class CleanUpPendingBuildsTable < ActiveRecord::Migration[6.0]
   def up
     return unless Gitlab.dev_or_test_env? || Gitlab.com?
 
-    each_batch_range('ci_pending_builds', of: BATCH_SIZE) do |min, max|
+    each_batch_range('ci_pending_builds', connection: connection, of: BATCH_SIZE) do |min, max|
       execute <<~SQL
         DELETE FROM ci_pending_builds
           USING ci_builds

db/post_migrate/20210610102413_migrate_protected_attribute_to_pending_builds.rb

@@ -8,7 +8,7 @@ class MigrateProtectedAttributeToPendingBuilds < ActiveRecord::Migration[6.1]
   def up
     return unless Gitlab.dev_or_test_env? || Gitlab.com?
 
-    each_batch_range('ci_pending_builds', of: 1000) do |min, max|
+    each_batch_range('ci_pending_builds', connection: connection, of: 1000) do |min, max|
       execute <<~SQL
         UPDATE ci_pending_builds
           SET protected = true

lib/gitlab/background_migration/backfill_integrations_type_new.rb

@@ -9,7 +9,7 @@ module Gitlab
       include Gitlab::Database::DynamicModelHelpers
 
       def perform(start_id, stop_id, batch_table, batch_column, sub_batch_size, pause_ms)
-        parent_batch_relation = define_batchable_model(batch_table)
+        parent_batch_relation = define_batchable_model(batch_table, connection: connection)
           .where(batch_column => start_id..stop_id)
 
         parent_batch_relation.each_batch(column: batch_column, of: sub_batch_size) do |sub_batch|

lib/gitlab/background_migration/backfill_user_namespace.rb

@@ -29,7 +29,7 @@ module Gitlab
       end
 
       def relation_scoped_to_range(source_table, source_key_column, start_id, stop_id)
-        define_batchable_model(source_table)
+        define_batchable_model(source_table, connection: connection)
           .where(source_key_column => start_id..stop_id)
           .where(type: nil)
       end

lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb

@@ -19,7 +19,7 @@ module Gitlab
         # batch_size - The size of the next batch
         # job_arguments - The migration job arguments
         def next_batch(table_name, column_name, batch_min_value:, batch_size:, job_arguments:)
-          model_class = define_batchable_model(table_name)
+          model_class = define_batchable_model(table_name, connection: ActiveRecord::Base.connection)
 
           quoted_column_name = model_class.connection.quote_column_name(column_name)
           relation = model_class.where("#{quoted_column_name} >= ?", batch_min_value)

lib/gitlab/background_migration/copy_column_using_background_migration_job.rb

@@ -57,7 +57,7 @@ module Gitlab
       end
 
       def relation_scoped_to_range(source_table, source_key_column, start_id, stop_id)
-        define_batchable_model(source_table).where(source_key_column => start_id..stop_id)
+        define_batchable_model(source_table, connection: connection).where(source_key_column => start_id..stop_id)
       end
 
       def column_assignment_clauses(copy_from, copy_to)

lib/gitlab/background_migration/delete_orphaned_deployments.rb

@@ -15,7 +15,7 @@ module Gitlab
       end
 
       def orphaned_deployments
-        define_batchable_model('deployments')
+        define_batchable_model('deployments', connection: ActiveRecord::Base.connection)
           .where('NOT EXISTS (SELECT 1 FROM environments WHERE deployments.environment_id = environments.id)')
       end
 

lib/gitlab/background_migration/remove_vulnerability_finding_links.rb

@@ -10,7 +10,9 @@ module Gitlab
       include Gitlab::Database::DynamicModelHelpers
 
       def perform(start_id, stop_id)
-        define_batchable_model('vulnerability_finding_links').where(id: start_id..stop_id).delete_all
+        define_batchable_model('vulnerability_finding_links', connection: ActiveRecord::Base.connection)
+          .where(id: start_id..stop_id)
+          .delete_all
       end
     end
   end

lib/gitlab/background_migration/update_timelogs_null_spent_at.rb

@@ -9,7 +9,9 @@ module Gitlab
       BATCH_SIZE = 100
 
       def perform(start_id, stop_id)
-        define_batchable_model('timelogs').where(spent_at: nil, id: start_id..stop_id).each_batch(of: 100) do |subbatch|
+        define_batchable_model('timelogs', connection: connection)
+            .where(spent_at: nil, id: start_id..stop_id)
+            .each_batch(of: 100) do |subbatch|
           batch_start, batch_end = subbatch.pluck('min(id), max(id)').first
 
           update_timelogs(batch_start, batch_end)
@@ -25,9 +27,12 @@ module Gitlab
         SQL
       end
 
-      def execute(sql)
+      def connection
         @connection ||= ::ActiveRecord::Base.connection
-        @connection.execute(sql)
+      end
+
+      def execute(sql)
+        connection.execute(sql)
       end
     end
   end

lib/gitlab/database/dynamic_model_helpers.rb

@@ -5,16 +5,19 @@ module Gitlab
     module DynamicModelHelpers
       BATCH_SIZE = 1_000
 
-      def define_batchable_model(table_name)
-        Class.new(ActiveRecord::Base) do
+      def define_batchable_model(table_name, connection:)
+        klass = Class.new(ActiveRecord::Base) do
           include EachBatch
 
           self.table_name = table_name
           self.inheritance_column = :_type_disabled
         end
+
+        klass.connection = connection
+        klass
       end
 
-      def each_batch(table_name, scope: ->(table) { table.all }, of: BATCH_SIZE)
+      def each_batch(table_name, connection:, scope: ->(table) { table.all }, of: BATCH_SIZE)
         if transaction_open?
           raise <<~MSG.squish
             each_batch should not run inside a transaction, you can disable
@@ -23,12 +26,12 @@ module Gitlab
           MSG
         end
 
-        scope.call(define_batchable_model(table_name))
+        scope.call(define_batchable_model(table_name, connection: connection))
           .each_batch(of: of) { |batch| yield batch }
       end
 
-      def each_batch_range(table_name, scope: ->(table) { table.all }, of: BATCH_SIZE)
-        each_batch(table_name, scope: scope, of: of) do |batch|
+      def each_batch_range(table_name, connection:, scope: ->(table) { table.all }, of: BATCH_SIZE)
+        each_batch(table_name, connection: connection, scope: scope, of: of) do |batch|
           yield batch.pluck('MIN(id), MAX(id)').first
         end
       end

lib/gitlab/database/migration_helpers.rb

@@ -9,6 +9,18 @@ module Gitlab
       include RenameTableHelpers
       include AsyncIndexes::MigrationHelpers
 
+      def define_batchable_model(table_name, connection: self.connection)
+        super(table_name, connection: connection)
+      end
+
+      def each_batch(table_name, connection: self.connection, **kwargs)
+        super(table_name, connection: connection, **kwargs)
+      end
+
+      def each_batch_range(table_name, connection: self.connection, **kwargs)
+        super(table_name, connection: connection, **kwargs)
+      end
+
       # https://www.postgresql.org/docs/current/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
       MAX_IDENTIFIER_NAME_LENGTH = 63
       DEFAULT_TIMESTAMP_COLUMNS = %i[created_at updated_at].freeze

lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb

@@ -49,7 +49,7 @@ module Gitlab
         end
 
         def relation_scoped_to_range(source_table, source_key_column, start_id, stop_id)
-          define_batchable_model(source_table)
+          define_batchable_model(source_table, connection: connection)
             .where(source_key_column => start_id..stop_id)
         end
 

lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb

@@ -5,7 +5,6 @@ module Gitlab
     module PartitioningMigrationHelpers
       module TableManagementHelpers
         include ::Gitlab::Database::SchemaHelpers
-        include ::Gitlab::Database::DynamicModelHelpers
         include ::Gitlab::Database::MigrationHelpers
         include ::Gitlab::Database::Migrations::BackgroundMigrationHelpers
 

spec/lib/gitlab/database/dynamic_model_helpers_spec.rb

@@ -4,10 +4,11 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::DynamicModelHelpers do
   let(:including_class) { Class.new.include(described_class) }
-  let(:table_name) { 'projects' }
+  let(:table_name) { Project.table_name }
+  let(:connection) { Project.connection }
 
   describe '#define_batchable_model' do
-    subject { including_class.new.define_batchable_model(table_name) }
+    subject { including_class.new.define_batchable_model(table_name, connection: connection) }
 
     it 'is an ActiveRecord model' do
       expect(subject.ancestors).to include(ActiveRecord::Base)
@@ -40,7 +41,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
 
       it 'iterates table in batches' do
         each_batch_size = ->(&block) do
-          subject.each_batch(table_name, of: 1) do |batch|
+          subject.each_batch(table_name, connection: connection, of: 1) do |batch|
             block.call(batch.size)
           end
         end
@@ -56,7 +57,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       end
 
       it 'raises an error' do
-        expect { subject.each_batch(table_name, of: 1) { |batch| batch.size } }
+        expect { subject.each_batch(table_name, connection: connection, of: 1) { |batch| batch.size } }
           .to raise_error(RuntimeError, /each_batch should not run inside a transaction/)
       end
     end
@@ -74,7 +75,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       end
 
       it 'iterates table in batch ranges' do
-        expect { |b| subject.each_batch_range(table_name, of: 1, &b) }
+        expect { |b| subject.each_batch_range(table_name, connection: connection, of: 1, &b) }
           .to yield_successive_args(
             [first_project.id, first_project.id],
             [second_project.id, second_project.id]
@@ -82,13 +83,13 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       end
 
       it 'yields only one batch if bigger than the table size' do
-        expect { |b| subject.each_batch_range(table_name, of: 2, &b) }
+        expect { |b| subject.each_batch_range(table_name, connection: connection, of: 2, &b) }
           .to yield_successive_args([first_project.id, second_project.id])
       end
 
       it 'makes it possible to apply a scope' do
         each_batch_limited = ->(&b) do
-          subject.each_batch_range(table_name, scope: ->(table) { table.limit(1) }, of: 1, &b)
+          subject.each_batch_range(table_name, connection: connection, scope: ->(table) { table.limit(1) }, of: 1, &b)
         end
 
         expect { |b| each_batch_limited.call(&b) }
@@ -102,7 +103,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       end
 
       it 'raises an error' do
-        expect { subject.each_batch_range(table_name, of: 1) { 1 } }
+        expect { subject.each_batch_range(table_name, connection: connection, of: 1) { 1 } }
           .to raise_error(RuntimeError, /each_batch should not run inside a transaction/)
       end
     end

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -14,6 +14,54 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     allow(model).to receive(:puts)
   end
 
+  describe 'overridden dynamic model helpers' do
+    let(:test_table) { '__test_batching_table' }
+
+    before do
+      model.connection.execute(<<~SQL)
+        CREATE TABLE #{test_table} (
+          id integer NOT NULL PRIMARY KEY,
+          name text NOT NULL
+        );
+
+        INSERT INTO #{test_table} (id, name)
+        VALUES (1, 'bob'), (2, 'mary'), (3, 'amy');
+      SQL
+    end
+
+    describe '#define_batchable_model' do
+      it 'defines a batchable model with the migration connection' do
+        expect(model.define_batchable_model(test_table).count).to eq(3)
+      end
+    end
+
+    describe '#each_batch' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+      end
+
+      it 'calls each_batch with the migration connection' do
+        each_batch_name = ->(&block) do
+          model.each_batch(test_table, of: 2) do |batch|
+            block.call(batch.pluck(:name))
+          end
+        end
+
+        expect { |b| each_batch_name.call(&b) }.to yield_successive_args(%w[bob mary], %w[amy])
+      end
+    end
+
+    describe '#each_batch_range' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+      end
+
+      it 'calls each_batch with the migration connection' do
+        expect { |b| model.each_batch_range(test_table, of: 2, &b) }.to yield_successive_args([1, 2], [3, 3])
+      end
+    end
+  end
+
   describe '#remove_timestamps' do
     it 'can remove the default timestamps' do
       Gitlab::Database::MigrationHelpers::DEFAULT_TIMESTAMP_COLUMNS.each do |column_name|