Add latest changes from gitlab-org/gitlab@master

3 files changed, 21 insertions, 54 deletions

diff --git a/lib/gitlab/database/background_migration_job.rb b/lib/gitlab/database/background_migration_job.rb
index c0e3016fd3d..5141dd05e4e 100644
--- a/lib/gitlab/database/background_migration_job.rb
+++ b/lib/gitlab/database/background_migration_job.rb
@@ -13,10 +13,6 @@ module Gitlab
         for_migration_class(class_name).where('arguments = ?', arguments.to_json) # rubocop:disable Rails/WhereEquals
       end
 
-      scope :for_partitioning_migration, -> (class_name, table_name) do
-        for_migration_class(class_name).where('arguments ->> 2 = ?', table_name)
-      end
-
       enum status: {
         pending: 0,
         succeeded: 1
diff --git a/lib/gitlab/database/load_balancing/connection_proxy.rb b/lib/gitlab/database/load_balancing/connection_proxy.rb
index 0d39b47dbba..02f14e020c1 100644
--- a/lib/gitlab/database/load_balancing/connection_proxy.rb
+++ b/lib/gitlab/database/load_balancing/connection_proxy.rb
@@ -32,7 +32,6 @@ module Gitlab
           select_one
           select_rows
           quote_column_name
-          schema_cache
         ).freeze
 
         # hosts - The hosts to use for load balancing.
@@ -63,6 +62,13 @@ module Gitlab
           end
         end
 
+        def schema_cache(*args, **kwargs, &block)
+          # Ignore primary stickiness for schema_cache queries and always use replicas
+          @load_balancer.read do |connection|
+            connection.public_send(:schema_cache, *args, **kwargs, &block)
+          end
+        end
+
         def transaction(*args, **kwargs, &block)
           if current_session.fallback_to_replicas_for_ambiguous_queries?
             track_read_only_transaction!
diff --git a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
index e3cf1298df6..8b49cb00bdf 100644
--- a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
@@ -17,12 +17,6 @@ module Gitlab
         BATCH_SIZE = 50_000
         SUB_BATCH_SIZE = 2_500
 
-        JobArguments = Struct.new(:start_id, :stop_id, :source_table_name, :partitioned_table_name, :source_column) do
-          def self.from_array(arguments)
-            self.new(*arguments)
-          end
-        end
-
         # Creates a partitioned copy of an existing table, using a RANGE partitioning strategy on a timestamp column.
         # One partition is created per month between the given `min_date` and `max_date`. Also installs a trigger on
         # the original table to copy writes into the partitioned table. To copy over historic data from before creation
@@ -124,7 +118,7 @@ module Gitlab
 
         # Cleanup a previously enqueued background migration to copy data into a partitioned table. This will not
         # prevent the enqueued jobs from executing, but instead cleans up information in the database used to track the
-        # state of the background migration. It should be safe to also remove the partitioned table even if the
+        # state of the batched background migration. It should be safe to also remove the partitioned table even if the
         # background jobs are still in-progress, as the absence of the table will cause them to safely exit.
         #
         # Example:
@@ -134,7 +128,10 @@ module Gitlab
         def cleanup_partitioning_data_migration(table_name)
           assert_table_is_allowed(table_name)
 
-          cleanup_migration_jobs(table_name)
+          partitioned_table_name = make_partitioned_table_name(table_name)
+          primary_key = connection.primary_key(table_name)
+
+          delete_batched_background_migration(MIGRATION, table_name, primary_key, [partitioned_table_name])
         end
 
         def create_hash_partitions(table_name, number_of_partitions)
@@ -154,11 +151,8 @@ module Gitlab
           end
         end
 
-        # Executes cleanup tasks from a previous BackgroundMigration to backfill a partitioned table by finishing
-        # pending jobs and performing a final data synchronization.
-        # This performs two steps:
-        #   1. Wait to finish any pending BackgroundMigration jobs that have not succeeded
-        #   2. Inline copy any missed rows from the original table to the partitioned table
+        # Executes jobs from previous BatchedBackgroundMigration to backfill the partitioned table by finishing
+        # pending jobs.
         #
         # **NOTE** Migrations using this method cannot be scheduled in the same release as the migration that
         # schedules the background migration using the `enqueue_partitioning_data_migration` helper, or else the
@@ -169,23 +163,21 @@ module Gitlab
         #   finalize_backfilling_partitioned_table :audit_events
         #
         def finalize_backfilling_partitioned_table(table_name)
-          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_dml_mode!
-
           assert_table_is_allowed(table_name)
-          assert_not_in_transaction_block(scope: ERROR_SCOPE)
 
           partitioned_table_name = make_partitioned_table_name(table_name)
+
           unless table_exists?(partitioned_table_name)
             raise "could not find partitioned table for #{table_name}, " \
               "this could indicate the previous partitioning migration has been rolled back."
           end
 
-          Gitlab::BackgroundMigration.steal(MIGRATION_CLASS_NAME) do |background_job|
-            JobArguments.from_array(background_job.args.second).source_table_name == table_name.to_s
-          end
-
-          primary_key = connection.primary_key(table_name)
-          copy_missed_records(table_name, partitioned_table_name, primary_key)
+          ensure_batched_background_migration_is_finished(
+            job_class_name: MIGRATION,
+            table_name: table_name,
+            column_name: connection.primary_key(table_name),
+            job_arguments: [partitioned_table_name]
+          )
 
           Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.with_suppressed do
             disable_statement_timeout do
@@ -456,33 +448,6 @@ module Gitlab
           create_trigger(table_name, trigger_name, function_name, fires: 'AFTER INSERT OR UPDATE OR DELETE')
         end
 
-        def cleanup_migration_jobs(table_name)
-          ::Gitlab::Database::BackgroundMigrationJob.for_partitioning_migration(MIGRATION_CLASS_NAME, table_name).delete_all
-        end
-
-        def copy_missed_records(source_table_name, partitioned_table_name, source_column)
-          backfill_table = BackfillPartitionedTable.new(connection: connection)
-
-          relation = ::Gitlab::Database::BackgroundMigrationJob.pending
-            .for_partitioning_migration(MIGRATION_CLASS_NAME, source_table_name)
-
-          relation.each_batch do |batch|
-            batch.each do |pending_migration_job|
-              job_arguments = JobArguments.from_array(pending_migration_job.arguments)
-              start_id = job_arguments.start_id
-              stop_id = job_arguments.stop_id
-
-              say("Backfilling data into partitioned table for ids from #{start_id} to #{stop_id}")
-              job_updated_count = backfill_table.perform(start_id, stop_id, source_table_name,
-                partitioned_table_name, source_column)
-
-              unless job_updated_count > 0
-                raise "failed to update tracking record for ids from #{start_id} to #{stop_id}"
-              end
-            end
-          end
-        end
-
         def replace_table(original_table_name, replacement_table_name, replaced_table_name, primary_key_name)
           replace_table = Gitlab::Database::Partitioning::ReplaceTable.new(connection,
               original_table_name.to_s, replacement_table_name, replaced_table_name, primary_key_name)