diff options
author | GitLab Bot <gitlab-bot@gitlab.com> | 2021-11-18 03:13:32 +0300 |
---|---|---|
committer | GitLab Bot <gitlab-bot@gitlab.com> | 2021-11-18 03:13:32 +0300 |
commit | d150848f7ed150d4f1887a4bff7cc4c423c55186 (patch) | |
tree | c43ccbe2e72615f7a60e647a4ce3ef2fa11dd35e /spec/support | |
parent | ed28831dfe5fd267a9506cb6642e50fbcb581a7f (diff) |
Add latest changes from gitlab-org/gitlab@master
Diffstat (limited to 'spec/support')
-rw-r--r-- | spec/support/database/prevent_cross_database_modification.rb | 147 |
1 files changed, 15 insertions, 132 deletions
diff --git a/spec/support/database/prevent_cross_database_modification.rb b/spec/support/database/prevent_cross_database_modification.rb index 09a87f92072..c509aecf9b8 100644 --- a/spec/support/database/prevent_cross_database_modification.rb +++ b/spec/support/database/prevent_cross_database_modification.rb @@ -1,148 +1,31 @@ # frozen_string_literal: true -module Database - module PreventCrossDatabaseModification - CrossDatabaseModificationAcrossUnsupportedTablesError = Class.new(StandardError) - - module GitlabDatabaseMixin - def allow_cross_database_modification_within_transaction(url:) - cross_database_context = Database::PreventCrossDatabaseModification.cross_database_context - return yield unless cross_database_context && cross_database_context[:enabled] - - transaction_tracker_enabled_was = cross_database_context[:enabled] - cross_database_context[:enabled] = false - - yield - ensure - cross_database_context[:enabled] = transaction_tracker_enabled_was if cross_database_context - end - end - - module SpecHelpers - def with_cross_database_modification_prevented - subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |name, start, finish, id, payload| - PreventCrossDatabaseModification.prevent_cross_database_modification!(payload[:connection], payload[:sql]) - end - - PreventCrossDatabaseModification.reset_cross_database_context! - PreventCrossDatabaseModification.cross_database_context.merge!(enabled: true, subscriber: subscriber) - - yield if block_given? - ensure - cleanup_with_cross_database_modification_prevented if block_given? - end - - def cleanup_with_cross_database_modification_prevented - if PreventCrossDatabaseModification.cross_database_context - ActiveSupport::Notifications.unsubscribe(PreventCrossDatabaseModification.cross_database_context[:subscriber]) - PreventCrossDatabaseModification.cross_database_context[:enabled] = false - end - end - end - - def self.cross_database_context - Thread.current[:transaction_tracker] - end - - def self.reset_cross_database_context! - Thread.current[:transaction_tracker] = initial_data - end - - def self.initial_data - { - enabled: false, - transaction_depth_by_db: Hash.new { |h, k| h[k] = 0 }, - modified_tables_by_db: Hash.new { |h, k| h[k] = Set.new } - } - end - - def self.prevent_cross_database_modification!(connection, sql) - return unless cross_database_context - return unless cross_database_context[:enabled] - - return if connection.pool.instance_of?(ActiveRecord::ConnectionAdapters::NullPool) - return if in_factory_bot_create? - - database = connection.pool.db_config.name - - if sql.start_with?('SAVEPOINT') - cross_database_context[:transaction_depth_by_db][database] += 1 - - return - elsif sql.start_with?('RELEASE SAVEPOINT', 'ROLLBACK TO SAVEPOINT') - cross_database_context[:transaction_depth_by_db][database] -= 1 - if cross_database_context[:transaction_depth_by_db][database] <= 0 - cross_database_context[:modified_tables_by_db][database].clear - end - - return - end - - return if cross_database_context[:transaction_depth_by_db].values.all?(&:zero?) - - # PgQuery might fail in some cases due to limited nesting: - # https://github.com/pganalyze/pg_query/issues/209 - parsed_query = PgQuery.parse(sql) - tables = sql.downcase.include?(' for update') ? parsed_query.tables : parsed_query.dml_tables - - # We have some code where plans and gitlab_subscriptions are lazily - # created and this causes lots of spec failures - # https://gitlab.com/gitlab-org/gitlab/-/issues/343394 - tables -= %w[plans gitlab_subscriptions] - - return if tables.empty? - - # All migrations will write to schema_migrations in the same transaction. - # It's safe to ignore this since schema_migrations exists in all - # databases - return if tables == ['schema_migrations'] - - cross_database_context[:modified_tables_by_db][database].merge(tables) - - all_tables = cross_database_context[:modified_tables_by_db].values.map(&:to_a).flatten - schemas = ::Gitlab::Database::GitlabSchema.table_schemas(all_tables) - - if schemas.many? - message = "Cross-database data modification of '#{schemas.to_a.join(", ")}' were detected within " \ - "a transaction modifying the '#{all_tables.to_a.join(", ")}' tables." \ - "Please refer to https://docs.gitlab.com/ee/development/database/multiple_databases.html#removing-cross-database-transactions for details on how to resolve this exception." - - if schemas.any? { |s| s.to_s.start_with?("undefined") } - message += " The gitlab_schema was undefined for one or more of the tables in this transaction. Any new tables must be added to spec/support/database/gitlab_schemas.yml ." - end - - raise Database::PreventCrossDatabaseModification::CrossDatabaseModificationAcrossUnsupportedTablesError, message - end - end - - # We ignore execution in the #create method from FactoryBot - # because it is not representative of real code we run in - # production. There are far too many false positives caused - # by instantiating objects in different `gitlab_schema` in a - # FactoryBot `create`. - def self.in_factory_bot_create? - caller_locations.any? { |l| l.path.end_with?('lib/factory_bot/evaluation.rb') && l.label == 'create' } - end - end +module PreventCrossDatabaseModificationSpecHelpers + delegate :with_cross_database_modification_prevented, + :allow_cross_database_modification_within_transaction, + to: :'::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification' end -Gitlab::Database.singleton_class.prepend( - Database::PreventCrossDatabaseModification::GitlabDatabaseMixin) - CROSS_DB_MODIFICATION_ALLOW_LIST = Set.new(YAML.load_file(File.join(__dir__, 'cross-database-modification-allowlist.yml'))).freeze RSpec.configure do |config| - config.include(::Database::PreventCrossDatabaseModification::SpecHelpers) + config.include(PreventCrossDatabaseModificationSpecHelpers) + + # By default allow cross-modifications as we want to observe only transactions + # within a specific block of execution which is defined be `before(:each)` and `after(:each)` + config.before(:all) do + ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.suppress = true + end # Using before and after blocks because the around block causes problems with the let_it_be # record creations. It makes an extra savepoint which breaks the transaction count logic. config.before do |example_file| - if CROSS_DB_MODIFICATION_ALLOW_LIST.exclude?(example_file.file_path_rerun_argument) - with_cross_database_modification_prevented - end + ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.suppress = + CROSS_DB_MODIFICATION_ALLOW_LIST.include?(example_file.file_path_rerun_argument) end + # Reset after execution to preferred state config.after do |example_file| - cleanup_with_cross_database_modification_prevented + ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.suppress = true end end |