Add allow list on callsite offenders

Add method to restore check when method ends

.cross-join-allowlist.yml

@@ -97,7 +97,6 @@
 - "./ee/spec/services/ci/minutes/refresh_cached_data_service_spec.rb"
 - "./ee/spec/services/ci/pipeline_creation/drop_not_runnable_builds_service_spec.rb"
 - "./ee/spec/services/ci/process_pipeline_service_spec.rb"
-- "./ee/spec/services/ci/register_job_service_spec.rb"
 - "./ee/spec/services/ci/retry_build_service_spec.rb"
 - "./ee/spec/services/ci/retry_pipeline_service_spec.rb"
 - "./ee/spec/services/ci/trigger_downstream_subscription_service_spec.rb"

app/models/ci/runner.rb

@@ -420,14 +420,18 @@ module Ci
     end
 
     def no_projects
-      if projects.any?
-        errors.add(:runner, 'cannot have projects assigned')
+      ::Gitlab::Database.allow_cross_joins_across_databases(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/338659') do
+        if projects.any?
+          errors.add(:runner, 'cannot have projects assigned')
+        end
       end
     end
 
     def no_groups
-      if groups.any?
-        errors.add(:runner, 'cannot have groups assigned')
+      ::Gitlab::Database.allow_cross_joins_across_databases(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/338659') do
+        if groups.any?
+          errors.add(:runner, 'cannot have groups assigned')
+        end
       end
     end
 

app/services/ci/register_job_service.rb

@@ -103,40 +103,42 @@ module Ci
 
     # rubocop: disable CodeReuse/ActiveRecord
     def each_build(params, &blk)
-      queue = ::Ci::Queue::BuildQueueService.new(runner)
-
-      builds = begin
-        if runner.instance_type?
-          queue.builds_for_shared_runner
-        elsif runner.group_type?
-          queue.builds_for_group_runner
-        else
-          queue.builds_for_project_runner
+      ::Gitlab::Database.allow_cross_joins_across_databases(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/339429') do
+        queue = ::Ci::Queue::BuildQueueService.new(runner)
+
+        builds = begin
+          if runner.instance_type?
+            queue.builds_for_shared_runner
+          elsif runner.group_type?
+            queue.builds_for_group_runner
+          else
+            queue.builds_for_project_runner
+          end
         end
-      end
 
-      if runner.ref_protected?
-        builds = queue.builds_for_protected_runner(builds)
-      end
+        if runner.ref_protected?
+          builds = queue.builds_for_protected_runner(builds)
+        end
 
-      # pick builds that does not have other tags than runner's one
-      builds = queue.builds_matching_tag_ids(builds, runner.tags.ids)
+        # pick builds that does not have other tags than runner's one
+        builds = queue.builds_matching_tag_ids(builds, runner.tags.ids)
 
-      # pick builds that have at least one tag
-      unless runner.run_untagged?
-        builds = queue.builds_with_any_tags(builds)
-      end
+        # pick builds that have at least one tag
+        unless runner.run_untagged?
+          builds = queue.builds_with_any_tags(builds)
+        end
 
-      # pick builds that older than specified age
-      if params.key?(:job_age)
-        builds = queue.builds_queued_before(builds, params[:job_age].seconds.ago)
-      end
+        # pick builds that older than specified age
+        if params.key?(:job_age)
+          builds = queue.builds_queued_before(builds, params[:job_age].seconds.ago)
+        end
 
-      build_ids = retrieve_queue(-> { queue.execute(builds) })
+        build_ids = retrieve_queue(-> { queue.execute(builds) })
 
-      @metrics.observe_queue_size(-> { build_ids.size }, @runner.runner_type)
+        @metrics.observe_queue_size(-> { build_ids.size }, @runner.runner_type)
 
-      build_ids.each { |build_id| yield Ci::Build.find(build_id) }
+        build_ids.each { |build_id| yield Ci::Build.find(build_id) }
+      end
     end
     # rubocop: enable CodeReuse/ActiveRecord
 

lib/gitlab/database.rb

@@ -145,6 +145,7 @@ module Gitlab
     def self.allow_cross_joins_across_databases(url:)
       # this method is implemented in:
       # spec/support/database/prevent_cross_joins.rb
+      yield
     end
 
     # This method will allow cross database modifications within the block

spec/support/database/prevent_cross_joins.rb

@@ -11,7 +11,7 @@
 #
 # class User
 #   def ci_owned_runners
-#     ::Gitlab::Database.allow_cross_joins_across_databases!(url: link-to-issue-url)
+#     ::Gitlab::Database.allow_cross_joins_across_databases(url: link-to-issue-url)
 #
 #     ...
 #   end
@@ -66,7 +66,10 @@ module Database
     module GitlabDatabaseMixin
       def allow_cross_joins_across_databases(url:)
         Thread.current[:allow_cross_joins_across_databases] = true
-        super
+
+        yield
+      ensure
+        Thread.current[:allow_cross_joins_across_databases] = false
       end
     end
   end

spec/support_specs/database/prevent_cross_joins_spec.rb

@@ -24,9 +24,7 @@ RSpec.describe Database::PreventCrossJoins do
 
       context 'when allow_cross_joins_across_databases is used' do
         it 'does not raise exception' do
-          Gitlab::Database.allow_cross_joins_across_databases(url: 'http://issue-url')
-
-          expect { main_and_ci_query }.not_to raise_error
+          expect { main_and_ci_query_allowlisted }.not_to raise_error
         end
       end
     end
@@ -34,6 +32,12 @@ RSpec.describe Database::PreventCrossJoins do
 
   private
 
+  def main_and_ci_query_allowlisted
+    Gitlab::Database.allow_cross_joins_across_databases(url: 'http://issue-url') do
+      main_and_ci_query
+    end
+  end
+
   def main_only_query
     Issue.joins(:project).last
   end