Add latest changes from gitlab-org/gitlab@master

3 files changed, 15 insertions, 23 deletions

diff --git a/lib/gitlab/database/migrations/batched_background_migration_helpers.rb b/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
index 079a14824c0..0ba7e1b0fd0 100644
--- a/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
+++ b/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
@@ -148,14 +148,6 @@ module Gitlab
               'your migration class.'
           end
 
-          database_name = Gitlab::Database.db_config_name(connection)
-
-          unless ActiveRecord::Base.configurations.primary?(database_name)
-            raise 'The `#finalize_background_migration` is currently not supported when running in decomposed database, ' \
-              'and this database is not `main:`. For more information visit: ' \
-              'https://docs.gitlab.com/ee/development/database/migrations_for_multiple_databases.html'
-          end
-
           Gitlab::Database::BackgroundMigration::BatchedMigration.reset_column_information
 
           migration = Gitlab::Database::BackgroundMigration::BatchedMigration.find_for_configuration(
@@ -163,7 +155,14 @@ module Gitlab
 
           raise 'Could not find batched background migration' if migration.nil?
 
-          Gitlab::Database::BackgroundMigration::BatchedMigrationRunner.finalize(job_class_name, table_name, column_name, job_arguments, connection: connection)
+          with_restored_connection_stack do |restored_connection|
+            Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.with_suppressed do
+              Gitlab::Database::BackgroundMigration::BatchedMigrationRunner.finalize(
+                job_class_name, table_name,
+                column_name, job_arguments,
+                connection: restored_connection)
+            end
+          end
         end
 
         # Deletes batched background migration for the given configuration.
diff --git a/lib/gitlab/database/migrations/reestablished_connection_stack.rb b/lib/gitlab/database/migrations/reestablished_connection_stack.rb
index d7cf482c32a..addc9d874af 100644
--- a/lib/gitlab/database/migrations/reestablished_connection_stack.rb
+++ b/lib/gitlab/database/migrations/reestablished_connection_stack.rb
@@ -17,7 +17,9 @@ module Gitlab
           original_handler = ActiveRecord::Base.connection_handler
 
           original_db_config = ActiveRecord::Base.connection_db_config
-          return yield if ActiveRecord::Base.configurations.primary?(original_db_config.name)
+          if ActiveRecord::Base.configurations.primary?(original_db_config.name)
+            return yield(ActiveRecord::Base.connection)
+          end
 
           # If the `ActiveRecord::Base` connection is different than `:main`
           # re-establish and configure `SharedModel` context accordingly
@@ -43,7 +45,7 @@ module Gitlab
           ActiveRecord::Base.establish_connection :main # rubocop:disable Database/EstablishConnection
 
           Gitlab::Database::SharedModel.using_connection(base_model.connection) do
-            yield
+            yield(base_model.connection)
           end
         ensure
           ActiveRecord::Base.connection_handler = original_handler
diff --git a/lib/gitlab/grape_logging/loggers/queue_duration_logger.rb b/lib/gitlab/grape_logging/loggers/queue_duration_logger.rb
index fe741a5bbe8..dde2bdd855e 100644
--- a/lib/gitlab/grape_logging/loggers/queue_duration_logger.rb
+++ b/lib/gitlab/grape_logging/loggers/queue_duration_logger.rb
@@ -6,21 +6,12 @@ module Gitlab
   module GrapeLogging
     module Loggers
       class QueueDurationLogger < ::GrapeLogging::Loggers::Base
-        attr_accessor :start_time
-
-        def before
-          @start_time = Time.now
-        end
-
         def parameters(request, _)
-          proxy_start = request.env['HTTP_GITLAB_WORKHORSE_PROXY_START'].presence
-
-          return {} unless proxy_start && start_time
+          duration_s = request.env[Gitlab::Middleware::RailsQueueDuration::GITLAB_RAILS_QUEUE_DURATION_KEY].presence
 
-          # Time in milliseconds since gitlab-workhorse started the request
-          duration = start_time.to_f * 1_000 - proxy_start.to_f / 1_000_000
+          return {} unless duration_s
 
-          { 'queue_duration_s': Gitlab::Utils.ms_to_round_sec(duration) }
+          { 'queue_duration_s': duration_s }
         end
       end
     end