78 lines
2.4 KiB
Ruby
78 lines
2.4 KiB
Ruby
|
# 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
|