diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb index 1b25d3c2090c20a5079640f5d2c98f69b945ce73..eb36011fcedde6d7720d07f39fed4da417ebb613 100644 --- a/lib/gitlab/database.rb +++ b/lib/gitlab/database.rb @@ -147,6 +147,11 @@ def self.allow_cross_joins_across_databases(url:) # spec/support/database/prevent_cross_joins.rb end + def self.allow_cross_database_modification_within_transaction(url:) + # this method is implemented in: + # spec/support/database/cross_database_modification_check.rb + end + def self.add_post_migrate_path_to_rails(force: false) return if ENV['SKIP_POST_DEPLOYMENT_MIGRATIONS'] && !force diff --git a/spec/support/database/ci_tables.rb b/spec/support/database/ci_tables.rb new file mode 100644 index 0000000000000000000000000000000000000000..99fc7ac2501fcd894981941d8364f22d54d9094e --- /dev/null +++ b/spec/support/database/ci_tables.rb @@ -0,0 +1,22 @@ +# frozen_string_literal: true + +# This module stores the CI-related database tables which are +# going to be moved to a separate database. +module Database + module CiTables + def self.include?(name) + ci_tables.include?(name) + end + + def self.ci_tables + @@ci_tables ||= Set.new.tap do |tables| # rubocop:disable Style/ClassVars + tables.merge(Ci::ApplicationRecord.descendants.map(&:table_name).compact) + + # It was decided that taggings/tags are best placed with CI + # https://gitlab.com/gitlab-org/gitlab/-/issues/333413 + tables.add('taggings') + tables.add('tags') + end + end + end +end diff --git a/spec/support/database/prevent_cross_database_modification.rb b/spec/support/database/prevent_cross_database_modification.rb new file mode 100644 index 0000000000000000000000000000000000000000..47202696dfd267bc49eba1401c818aefea21eb1c --- /dev/null +++ b/spec/support/database/prevent_cross_database_modification.rb @@ -0,0 +1,120 @@ +# frozen_string_literal: true + +module Database + module PreventCrossDatabaseModification + CrossDatabaseModificationAcrossUnsupportedTablesError = Class.new(StandardError) + + module GitlabDatabaseMixin + def allow_cross_database_modification_within_transaction(url:) + return yield unless Thread.current[:transaction_tracker] + + cross_database_context = Database::PreventCrossDatabaseModification.cross_database_context + return yield unless 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 Thread.current[:transaction_tracker] + 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 + ActiveSupport::Notifications.unsubscribe(PreventCrossDatabaseModification.cross_database_context[:subscriber]) + PreventCrossDatabaseModification.cross_database_context[:enabled] = false + end + end + + def self.cross_database_context + Thread.current[:transaction_tracker] ||= initial_data + 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[:enabled] + + 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?) + + tables = PgQuery.parse(sql).dml_tables + + return if tables.empty? + + cross_database_context[:modified_tables_by_db][database].merge(tables) + + ci_table_referenced = false + main_table_referenced = false + all_tables = cross_database_context[:modified_tables_by_db].values.map(&:to_a).flatten + all_tables.each do |table| + if Database::CiTables.include?(table) + ci_table_referenced = true + else + main_table_referenced = true + end + end + + if ci_table_referenced && main_table_referenced + raise Database::PreventCrossDatabaseModification::CrossDatabaseModificationAcrossUnsupportedTablesError, + "Cross-database data modification queries (CI and Main) were detected within " \ + "a transaction '#{all_tables.join(", ")}' discovered" + end + end + end +end + +Gitlab::Database.singleton_class.prepend( + Database::PreventCrossDatabaseModification::GitlabDatabaseMixin) + +RSpec.configure do |config| + config.include(::Database::PreventCrossDatabaseModification::SpecHelpers) + + # 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(:each, :prevent_cross_database_modification) do + with_cross_database_modification_prevented + end + + config.after(:each, :prevent_cross_database_modification) do + cleanup_with_cross_database_modification_prevented + end +end diff --git a/spec/support/database/prevent_cross_joins.rb b/spec/support/database/prevent_cross_joins.rb index c49b165c6c3d7fad08edcc8ad97b3bbc65b8d4fa..789721ccd38a1f80a40918502635f782be02c84d 100644 --- a/spec/support/database/prevent_cross_joins.rb +++ b/spec/support/database/prevent_cross_joins.rb @@ -37,23 +37,8 @@ def self.validate_cross_joins!(sql) # 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| ci_table_name?(table) } || - tables.none? { |table| ci_table_name?(table) } - end - - def self.ci_table_name?(name) - ci_tables.include?(name) - end - - def self.ci_tables - @@ci_tables ||= Set.new.tap do |tables| # rubocop:disable Style/ClassVars - tables.merge(Ci::ApplicationRecord.descendants.map(&:table_name).compact) - - # It was decided that taggings/tags are best placed with CI - # https://gitlab.com/gitlab-org/gitlab/-/issues/333413 - tables.add('taggings') - tables.add('tags') - end + tables.all? { |table| CiTables.include?(table) } || + tables.none? { |table| CiTables.include?(table) } end module SpecHelpers diff --git a/spec/support_specs/database/prevent_cross_database_modification_spec.rb b/spec/support_specs/database/prevent_cross_database_modification_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..4fd55d59db0b7e364416a4e825984953c6151a88 --- /dev/null +++ b/spec/support_specs/database/prevent_cross_database_modification_spec.rb @@ -0,0 +1,144 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe 'Database::PreventCrossDatabaseModification' do + let_it_be(:pipeline, refind: true) { create(:ci_pipeline) } + let_it_be(:project, refind: true) { create(:project) } + + shared_examples 'succeessful examples' do + context 'outside transaction' do + it { expect { run_queries }.not_to raise_error } + end + + context 'within transaction' do + it do + Project.transaction do + expect { run_queries }.not_to raise_error + end + end + end + + context 'within nested transaction' do + it do + Project.transaction(requires_new: true) do + Project.transaction(requires_new: true) do + expect { run_queries }.not_to raise_error + end + end + end + end + end + + context 'when CI and other tables are read in a transaction' do + def run_queries + pipeline.reload + project.reload + end + + include_examples 'succeessful examples' + end + + context 'when only CI data is modified' do + def run_queries + pipeline.touch + project.reload + end + + include_examples 'succeessful examples' + end + + context 'when other data is modified' do + def run_queries + pipeline.reload + project.touch + end + + include_examples 'succeessful examples' + end + + describe 'with_cross_database_modification_prevented block' do + it 'raises error when CI and other data is modified' do + expect do + with_cross_database_modification_prevented do + Project.transaction do + project.touch + pipeline.touch + end + end + end.to raise_error /Cross-database data modification queries/ + end + end + + context 'when running tests with prevent_cross_database_modification', :prevent_cross_database_modification do + context 'when both CI and other data is modified' do + def run_queries + project.touch + pipeline.touch + end + + context 'outside transaction' do + it { expect { run_queries }.not_to raise_error } + end + + context 'when data modification happens in a transaction' do + it 'raises error' do + Project.transaction do + expect { run_queries }.to raise_error /Cross-database data modification queries/ + end + end + + context 'when data modification happens in nested transactions' do + it 'raises error' do + Project.transaction(requires_new: true) do + project.touch + Project.transaction(requires_new: true) do + expect { pipeline.touch }.to raise_error /Cross-database data modification queries/ + end + end + end + end + end + end + + context 'when CI association is modified through project' do + def run_queries + project.variables.build(key: 'a', value: 'v') + project.save! + end + + include_examples 'succeessful examples' + end + + describe '#allow_cross_database_modification_within_transaction' do + it 'skips raising error' do + expect do + Gitlab::Database.allow_cross_database_modification_within_transaction(url: 'gitlab-issue') do + Project.transaction do + pipeline.touch + project.touch + end + end + end.not_to raise_error + end + + it 'raises error when complex factories are built referencing both databases' do + expect do + ApplicationRecord.transaction do + create(:ci_pipeline) + end + end.to raise_error /Cross-database data modification queries/ + end + + it 'skips raising error on factory creation' do + expect do + Gitlab::Database.allow_cross_database_modification_within_transaction(url: 'gitlab-issue') do + ApplicationRecord.transaction do + create(:ci_pipeline) + end + end + end.not_to raise_error + end + end + end +end