1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
# 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
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)
# 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
end
config.after do |example_file|
cleanup_with_cross_database_modification_prevented
end
end
|
