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
149
150
151
152
153
154
155
156
157
158
159
160
161
|
# frozen_string_literal: true
module Database
module MultipleDatabasesHelpers
EXTRA_DBS = ::Gitlab::Database.all_database_names.map(&:to_sym) - [:main]
def database_exists?(database_name)
::Gitlab::Database.has_database?(database_name)
end
def skip_if_shared_database(database_name)
skip "Skipping because #{database_name} is shared or doesn't not exist" unless database_exists?(database_name)
end
def skip_if_database_exists(database_name)
skip "Skipping because database #{database_name} exists" if database_exists?(database_name)
end
def execute_on_each_database(query, databases: %I[main ci])
databases = databases.select { |database_name| database_exists?(database_name) }
Gitlab::Database::EachDatabase.each_connection(only: databases, include_shared: false) do |connection, _|
next unless Gitlab::Database.gitlab_schemas_for_connection(connection).include?(:gitlab_shared)
connection.execute(query)
end
end
def skip_if_multiple_databases_not_setup(*databases)
unless (databases - EXTRA_DBS).empty?
raise "Unsupported database in #{databases}. It must be one of #{EXTRA_DBS}."
end
databases = EXTRA_DBS if databases.empty?
return if databases.any? { |db| Gitlab::Database.has_config?(db) }
skip "Skipping because none of the extra databases #{databases} are setup"
end
def skip_if_multiple_databases_are_setup(*databases)
unless (databases - EXTRA_DBS).empty?
raise "Unsupported database in #{databases}. It must be one of #{EXTRA_DBS}."
end
databases = EXTRA_DBS if databases.empty?
return if databases.none? { |db| Gitlab::Database.has_config?(db) }
skip "Skipping because some of the extra databases #{databases} are setup"
end
def reconfigure_db_connection(name: nil, config_hash: {}, model: ActiveRecord::Base, config_model: nil)
db_config = (config_model || model).connection_db_config
new_db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
db_config.env_name,
name ? name.to_s : db_config.name,
db_config.configuration_hash.merge(config_hash)
)
model.establish_connection(new_db_config)
end
def ensure_schema_and_empty_tables
# Ensure all schemas for both databases are migrated back
Gitlab::Database.database_base_models.each do |_, base_model|
with_reestablished_active_record_base do
reconfigure_db_connection(
model: ActiveRecord::Base,
config_model: base_model
)
# Delete after migrating so that rows created during migration don't impact other
# specs (for example, async foreign key creation rows)
schema_migrate_up!
delete_from_all_tables!(except: deletion_except_tables)
end
end
# ActiveRecord::Base.clear_all_connections! disconnects and clears attribute methods
# Force a refresh to avoid schema failures.
reset_column_in_all_models
refresh_attribute_methods
end
# The usage of this method switches temporarily used `connection_handler`
# allowing full manipulation of ActiveRecord::Base connections without
# having side effects like:
# - misaligned transactions since this is managed by `TestProf::BeforeAll::Adapters::ActiveRecord`
# - removal of primary connections
#
# The execution within a block ensures safe cleanup of all allocated resources.
#
def with_reestablished_active_record_base(reconnect: true)
connection_classes = ActiveRecord::Base
.connection_handler
.connection_pool_names
.map(&:constantize)
connection_classes.delete(ActiveRecord::PendingMigrationConnection) if ::Gitlab.next_rails?
connection_class_to_config = connection_classes.index_with(&:connection_db_config)
original_handler = ActiveRecord::Base.connection_handler
new_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
ActiveRecord::Base.connection_handler = new_handler
connection_class_to_config.each { |klass, db_config| klass.establish_connection(db_config) } if reconnect
yield
ensure
ActiveRecord::Base.connection_handler = original_handler
new_handler&.clear_all_connections!
end
def with_db_configs(test: test_config)
current_configurations = ActiveRecord::Base.configurations
ActiveRecord::Base.configurations = { test: test_config }
yield
ensure
ActiveRecord::Base.configurations = current_configurations
end
def with_added_ci_connection
if Gitlab::Database.has_config?(:ci)
# No need to add a ci: connection if we already have one
yield
else
with_reestablished_active_record_base(reconnect: true) do
reconfigure_db_connection(
name: :ci,
model: Ci::ApplicationRecord,
config_model: ActiveRecord::Base
)
yield
# Cleanup connection_specification_name for Ci::ApplicationRecord
Ci::ApplicationRecord.remove_connection
end
end
end
end
module ActiveRecordBaseEstablishConnection
def establish_connection(*args)
# rubocop:disable Database/MultipleDatabases
if connected? &&
connection&.transaction_open? &&
ActiveRecord::Base.connection_handler == ActiveRecord::Base.default_connection_handler
raise "Cannot re-establish '#{self}.establish_connection' within an open transaction " \
"(#{connection&.open_transactions.to_i}). Use `with_reestablished_active_record_base` " \
"instead or add `:reestablished_active_record_base` to rspec context."
end
# rubocop:enable Database/MultipleDatabases
super
end
end
end
ActiveRecord::Base.singleton_class.prepend(::Database::ActiveRecordBaseEstablishConnection) # rubocop:disable Database/MultipleDatabases
|
