1 files changed, 39 insertions, 18 deletions

diff --git a/spec/support/database/prevent_cross_joins.rb b/spec/support/database/prevent_cross_joins.rb
index 789721ccd38..4b78aa9014c 100644
--- a/spec/support/database/prevent_cross_joins.rb
+++ b/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
@@ -21,33 +21,43 @@ module Database
   module PreventCrossJoins
     CrossJoinAcrossUnsupportedTablesError = Class.new(StandardError)
 
+    ALLOW_THREAD_KEY = :allow_cross_joins_across_databases
+
     def self.validate_cross_joins!(sql)
-      return if Thread.current[:allow_cross_joins_across_databases]
+      return if Thread.current[ALLOW_THREAD_KEY]
+
+      # Allow spec/support/database_cleaner.rb queries to disable/enable triggers for many tables
+      # See https://gitlab.com/gitlab-org/gitlab/-/issues/339396
+      return if sql.include?("DISABLE TRIGGER") || sql.include?("ENABLE TRIGGER")
 
       # PgQuery might fail in some cases due to limited nesting:
       # https://github.com/pganalyze/pg_query/issues/209
-      tables = PgQuery.parse(sql).tables
+      #
+      # 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
 
-      unless only_ci_or_only_main?(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(", ")}' discovered " \
-          "when executing query '#{sql}'"
+          "Unsupported cross-join across '#{tables.join(", ")}' modifying '#{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
 
-    # Returns true if a set includes only CI tables, or includes only non-CI tables
-    def self.only_ci_or_only_main?(tables)
-      tables.all? { |table| CiTables.include?(table) } ||
-        tables.none? { |table| CiTables.include?(table) }
-    end
-
     module SpecHelpers
       def with_cross_joins_prevented
         subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event|
           ::Database::PreventCrossJoins.validate_cross_joins!(event.payload[:sql])
         end
 
-        Thread.current[:allow_cross_joins_across_databases] = false
+        Thread.current[ALLOW_THREAD_KEY] = false
 
         yield
       ensure
@@ -57,8 +67,12 @@ module Database
 
     module GitlabDatabaseMixin
       def allow_cross_joins_across_databases(url:)
-        Thread.current[:allow_cross_joins_across_databases] = true
-        super
+        old_value = Thread.current[ALLOW_THREAD_KEY]
+        Thread.current[ALLOW_THREAD_KEY] = true
+
+        yield
+      ensure
+        Thread.current[ALLOW_THREAD_KEY] = old_value
       end
     end
   end
@@ -67,11 +81,18 @@ end
 Gitlab::Database.singleton_class.prepend(
   Database::PreventCrossJoins::GitlabDatabaseMixin)
 
+ALLOW_LIST = Set.new(YAML.load_file(File.join(__dir__, 'cross-join-allowlist.yml'))).freeze
+
 RSpec.configure do |config|
   config.include(::Database::PreventCrossJoins::SpecHelpers)
 
-  # TODO: remove `:prevent_cross_joins` to enable the check by default
-  config.around(:each, :prevent_cross_joins) do |example|
-    with_cross_joins_prevented { example.run }
+  config.around do |example|
+    Thread.current[:has_cross_join_exception] = false
+
+    if ALLOW_LIST.include?(example.file_path)
+      example.run
+    else
+      with_cross_joins_prevented { example.run }
+    end
   end
 end