Commit 73bd174d authored by Dylan Griffith's avatar Dylan Griffith

Merge branch 'helper-to-prevent-multi-database-data-modification-within-transaction' into 'master'

Utility for detecting cross DB data modification

See merge request gitlab-org/gitlab!67213
parents b2a83a36 0415725d
......@@ -147,6 +147,11 @@ module Gitlab
# 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
......
# 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
# 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
......@@ -37,23 +37,8 @@ module Database
# 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
......
# 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
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment