1 files changed, 17 insertions, 10 deletions
diff --git a/spec/support/database/prevent_cross_joins.rb b/spec/support/database/prevent_cross_joins.rb
index 4b78aa9014c..f5ed2a8f22e 100644
--- a/spec/support/database/prevent_cross_joins.rb
+++ b/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
@@ -32,21 +33,14 @@ module Database
 
       # PgQuery might fail in some cases due to limited nesting:
       # https://github.com/pganalyze/pg_query/issues/209
-      #
-      # Also, we disable GC while parsing because of https://github.com/pganalyze/pg_query/issues/226
-      begin
-        GC.disable
-        tables = PgQuery.parse(sql).tables
-      ensure
-        GC.enable
-      end
+      tables = PgQuery.parse(sql).tables
 
       schemas = Database::GitlabSchema.table_schemas(tables)
 
       if schemas.include?(:gitlab_ci) && schemas.include?(:gitlab_main)
         Thread.current[:has_cross_join_exception] = true
         raise CrossJoinAcrossUnsupportedTablesError,
-          "Unsupported cross-join across '#{tables.join(", ")}' modifying '#{schemas.to_a.join(", ")}' discovered " \
+          "Unsupported cross-join across '#{tables.join(", ")}' querying '#{schemas.to_a.join(", ")}' discovered " \
           "when executing query '#{sql}'. Please refer to https://docs.gitlab.com/ee/development/database/multiple_databases.html#removing-joins-between-ci_-and-non-ci_-tables for details on how to resolve this exception."
       end
     end
@@ -63,6 +57,10 @@ module Database
       ensure
         ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
       end
+
+      def allow_cross_joins_across_databases(url:, &block)
+        ::Gitlab::Database.allow_cross_joins_across_databases(url: url, &block)
+      end
     end
 
     module GitlabDatabaseMixin
@@ -75,12 +73,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|