1 files changed, 77 insertions, 0 deletions
diff --git a/spec/support/database/prevent_cross_joins.rb b/spec/support/database/prevent_cross_joins.rb
new file mode 100644
index 00000000000..789721ccd38
--- /dev/null
+++ b/spec/support/database/prevent_cross_joins.rb
@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+# This module tries to discover and prevent cross-joins across tables
+# This will forbid usage of tables between CI and main database
+# on a same query unless explicitly allowed by. This will change execution
+# from a given point to allow cross-joins. The state will be cleared
+# on a next test run.
+#
+# This method should be used to mark METHOD introducing cross-join
+# not a test using the cross-join.
+#
+# class User
+#   def ci_owned_runners
+#     ::Gitlab::Database.allow_cross_joins_across_databases!(url: link-to-issue-url)
+#
+#     ...
+#   end
+# end
+
+module Database
+  module PreventCrossJoins
+    CrossJoinAcrossUnsupportedTablesError = Class.new(StandardError)
+
+    def self.validate_cross_joins!(sql)
+      return if Thread.current[:allow_cross_joins_across_databases]
+
+      # PgQuery might fail in some cases due to limited nesting:
+      # https://github.com/pganalyze/pg_query/issues/209
+      tables = PgQuery.parse(sql).tables
+
+      unless only_ci_or_only_main?(tables)
+        raise CrossJoinAcrossUnsupportedTablesError,
+          "Unsupported cross-join across '#{tables.join(", ")}' discovered " \
+          "when executing query '#{sql}'"
+      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
+
+        yield
+      ensure
+        ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+      end
+    end
+
+    module GitlabDatabaseMixin
+      def allow_cross_joins_across_databases(url:)
+        Thread.current[:allow_cross_joins_across_databases] = true
+        super
+      end
+    end
+  end
+end
+
+Gitlab::Database.singleton_class.prepend(
+  Database::PreventCrossJoins::GitlabDatabaseMixin)
+
+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 }
+  end
+end