Add `ActiveRecord::Relation#allow_cross_joins_across_databases` method

This uses `#annotate` to indicate that given relation is exempt
from cross-join verification check.

config/initializers/00_active_record_disable_joins.rb

+# frozen_string_literal: true
+
+module ActiveRecordRelationAllowCrossJoins
+  def allow_cross_joins_across_databases(url:)
+    # this method is implemented in:
+    # spec/support/database/prevent_cross_joins.rb
+    self
+  end
+end
+
+ActiveRecord::Relation.prepend(ActiveRecordRelationAllowCrossJoins)

doc/development/database/multiple_databases.md

@@ -394,7 +394,8 @@ You can see a real example of using this method for fixing a cross-join in
 #### Allowlist for existing cross-joins
 
 A cross-join across databases can be explicitly allowed by wrapping the code in the
-`::Gitlab::Database.allow_cross_joins_across_databases` helper method.
+`::Gitlab::Database.allow_cross_joins_across_databases` helper method. Alternative
+way is to mark a given relation as `relation.allow_cross_joins_across_databases`.
 
 This method should only be used:
 
@@ -405,11 +406,22 @@ This method should only be used:
 The `allow_cross_joins_across_databases` helper method can be used as follows:
 
 ```ruby
+# Scope the block executing a object from database
 ::Gitlab::Database.allow_cross_joins_across_databases(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/336590') do
   subject.perform(1, 4)
 end
 ```
 
+```ruby
+# Mark a relation as allowed to cross-join databases
+def find_actual_head_pipeline
+  all_pipelines
+    .allow_cross_joins_across_databases(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/336891')
+    .for_sha_or_source_sha(diff_head_sha)
+    .first
+end
+```
+
 The `url` parameter should point to an issue with a milestone for when we intend
 to fix the cross-join. If the cross-join is being used in a migration, we do not
 need to fix the code. See <https://gitlab.com/gitlab-org/gitlab/-/issues/340017>

spec/support/database/prevent_cross_joins.rb

@@ -22,9 +22,10 @@ module Database
     CrossJoinAcrossUnsupportedTablesError = Class.new(StandardError)
 
     ALLOW_THREAD_KEY = :allow_cross_joins_across_databases
+    ALLOW_ANNOTATE_KEY = ALLOW_THREAD_KEY.to_s.freeze
 
     def self.validate_cross_joins!(sql)
-      return if Thread.current[ALLOW_THREAD_KEY]
+      return if Thread.current[ALLOW_THREAD_KEY] || sql.include?(ALLOW_ANNOTATE_KEY)
 
       # Allow spec/support/database_cleaner.rb queries to disable/enable triggers for many tables
       # See https://gitlab.com/gitlab-org/gitlab/-/issues/339396
@@ -75,12 +76,21 @@ module Database
         Thread.current[ALLOW_THREAD_KEY] = old_value
       end
     end
+
+    module ActiveRecordRelationMixin
+      def allow_cross_joins_across_databases(url:)
+        super.annotate(ALLOW_ANNOTATE_KEY)
+      end
+    end
   end
 end
 
 Gitlab::Database.singleton_class.prepend(
   Database::PreventCrossJoins::GitlabDatabaseMixin)
 
+ActiveRecord::Relation.prepend(
+  Database::PreventCrossJoins::ActiveRecordRelationMixin)
+
 ALLOW_LIST = Set.new(YAML.load_file(File.join(__dir__, 'cross-join-allowlist.yml'))).freeze
 
 RSpec.configure do |config|

spec/support_specs/database/prevent_cross_joins_spec.rb

@@ -22,6 +22,12 @@ RSpec.describe Database::PreventCrossJoins do
           described_class::CrossJoinAcrossUnsupportedTablesError)
       end
 
+      context 'when annotation is used' do
+        it 'does not raise exception' do
+          expect { main_and_ci_allowed_via_annotate }.not_to raise_error
+        end
+      end
+
       context 'when allow_cross_joins_across_databases is used' do
         it 'does not raise exception' do
           expect { main_and_ci_query_allowlisted }.not_to raise_error
@@ -52,6 +58,12 @@ RSpec.describe Database::PreventCrossJoins do
     end
   end
 
+  def main_and_ci_allowed_via_annotate
+    main_and_ci_query do |relation|
+      relation.allow_cross_joins_across_databases(url: 'http://issue-url')
+    end
+  end
+
   def main_only_query
     Issue.joins(:project).last
   end
@@ -61,6 +73,8 @@ RSpec.describe Database::PreventCrossJoins do
   end
 
   def main_and_ci_query
-    Ci::Build.joins(:project).last
+    relation = Ci::Build.joins(:project)
+    relation = yield(relation) if block_given?
+    relation.last
   end
 end