Commit 6f706bb0 authored by Adam Hegyi's avatar Adam Hegyi

Merge branch 'fix-establish_connection-in-specs' into 'master'

Add `#with_reestablished_active_record_base` to allow usage of `establish_connection`

See merge request gitlab-org/gitlab!70606
parents 4c93874c a049dd5b
......@@ -2,7 +2,7 @@
require 'spec_helper'
RSpec.describe Gitlab::Geo::DatabaseTasks do
RSpec.describe Gitlab::Geo::DatabaseTasks, :reestablished_active_record_base do
let(:schema_file) { Rails.root.join('tmp', 'tests', 'geo_schema.rb').to_s }
subject { described_class }
......
......@@ -5,7 +5,7 @@ require 'spec_helper'
RSpec.describe SystemCheck::Geo::GeoDatabaseConfiguredCheck, :silence_stdout do
subject { described_class.new }
describe '#multi_check' do
describe '#multi_check', :reestablished_active_record_base do
it "checks database configuration" do
stub_configuration_check(false)
......
......@@ -235,7 +235,7 @@ module Gitlab
@configuration.model.connection_specification_name,
role: ActiveRecord::Base.writing_role,
shard: ActiveRecord::Base.default_shard
)
) || raise(::ActiveRecord::ConnectionNotEstablished)
end
private
......
......@@ -2,19 +2,11 @@
require 'spec_helper'
RSpec.describe 'Database config initializer' do
RSpec.describe 'Database config initializer', :reestablished_active_record_base do
subject do
load Rails.root.join('config/initializers/database_config.rb')
end
around do |example|
original_config = ActiveRecord::Base.connection_db_config
example.run
ActiveRecord::Base.establish_connection(original_config)
end
before do
allow(Gitlab::Runtime).to receive(:max_threads).and_return(max_threads)
end
......
......@@ -91,45 +91,38 @@ RSpec.describe Gitlab::Database::BulkUpdate do
.to eq(['MR a', 'Issue a', 'Issue b'])
end
shared_examples 'basic functionality' do
it 'sets multiple values' do
create_default(:user)
create_default(:project)
i_a, i_b = create_list(:issue, 2)
context 'validates prepared_statements support', :reestablished_active_record_base do
using RSpec::Parameterized::TableSyntax
mapping = {
i_a => { title: 'Issue a' },
i_b => { title: 'Issue b' }
}
where(:prepared_statements) do
[false, true]
end
described_class.execute(%i[title], mapping)
before do
configuration_hash = ActiveRecord::Base.connection_db_config.configuration_hash
expect([i_a, i_b].map { |x| x.reset.title })
.to eq(['Issue a', 'Issue b'])
ActiveRecord::Base.establish_connection(
configuration_hash.merge(prepared_statements: prepared_statements)
)
end
end
include_examples 'basic functionality'
with_them do
it 'sets multiple values' do
create_default(:user)
create_default(:project)
context 'when prepared statements are configured differently to the normal test environment' do
before do
klass = Class.new(ActiveRecord::Base) do
def self.abstract_class?
true # So it gets its own connection
end
end
i_a, i_b = create_list(:issue, 2)
stub_const('ActiveRecordBasePreparedStatementsInverted', klass)
mapping = {
i_a => { title: 'Issue a' },
i_b => { title: 'Issue b' }
}
c = ActiveRecord::Base.retrieve_connection.instance_variable_get(:@config)
inverted = c.merge(prepared_statements: !ActiveRecord::Base.connection.prepared_statements)
ActiveRecordBasePreparedStatementsInverted.establish_connection(inverted)
described_class.execute(%i[title], mapping)
allow(ActiveRecord::Base).to receive(:connection_specification_name)
.and_return(ActiveRecordBasePreparedStatementsInverted.connection_specification_name)
expect([i_a, i_b].map { |x| x.reset.title })
.to eq(['Issue a', 'Issue b'])
end
end
include_examples 'basic functionality'
end
end
......@@ -126,15 +126,7 @@ RSpec.describe Gitlab::Database::Connection do
end
end
describe '#disable_prepared_statements' do
around do |example|
original_config = connection.scope.connection.pool.db_config
example.run
connection.scope.establish_connection(original_config)
end
describe '#disable_prepared_statements', :reestablished_active_record_base do
it 'disables prepared statements' do
connection.scope.establish_connection(
::Gitlab::Database.main.config.merge(prepared_statements: true)
......
......@@ -2,7 +2,7 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::PostgresqlAdapter::ForceDisconnectableMixin do
RSpec.describe Gitlab::Database::PostgresqlAdapter::ForceDisconnectableMixin, :reestablished_active_record_base do
describe 'checking in a connection to the pool' do
let(:model) do
Class.new(ActiveRecord::Base) do
......
......@@ -23,7 +23,7 @@ RSpec.describe Gitlab::Database::SchemaMigrations::Context do
end
end
context 'multiple databases' do
context 'multiple databases', :reestablished_active_record_base do
let(:connection_class) do
Class.new(::ApplicationRecord) do
self.abstract_class = true
......@@ -34,8 +34,6 @@ RSpec.describe Gitlab::Database::SchemaMigrations::Context do
end
end
let(:configuration_overrides) { {} }
before do
connection_class.establish_connection(
ActiveRecord::Base
......
......@@ -5,5 +5,57 @@ module Database
def skip_if_multiple_databases_not_setup
skip 'Skipping because multiple databases not set up' unless Gitlab::Database.has_config?(:ci)
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 `BeforeAllAdapter`
# - removal of primary connections
#
# The execution within a block ensures safe cleanup of all allocated resources.
#
# rubocop:disable Database/MultipleDatabases
def with_reestablished_active_record_base(reconnect: true)
connection_classes = ActiveRecord::Base.connection_handler.connection_pool_names.map(&:constantize).to_h do |klass|
[klass, klass.connection_db_config]
end
original_handler = ActiveRecord::Base.connection_handler
new_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
ActiveRecord::Base.connection_handler = new_handler
if reconnect
connection_classes.each { |klass, db_config| klass.establish_connection(db_config) }
end
yield
ensure
ActiveRecord::Base.connection_handler = original_handler
new_handler&.clear_all_connections!
end
# rubocop:enable Database/MultipleDatabases
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
RSpec.configure do |config|
config.around(:each, :reestablished_active_record_base) do |example|
with_reestablished_active_record_base(reconnect: example.metadata.fetch(:reconnect, true)) do
example.run
end
end
end
ActiveRecord::Base.singleton_class.prepend(::Database::ActiveRecordBaseEstablishConnection) # rubocop:disable Database/MultipleDatabases
# frozen_string_literal: true
require 'spec_helper'
RSpec.describe 'Database::MultipleDatabases' do
describe '.with_reestablished_active_record_base' do
context 'when doing establish_connection' do
context 'on ActiveRecord::Base' do
it 'raises exception' do
expect { ActiveRecord::Base.establish_connection(:main) }.to raise_error /Cannot re-establish/
end
context 'when using with_reestablished_active_record_base' do
it 'does not raise exception' do
with_reestablished_active_record_base do
expect { ActiveRecord::Base.establish_connection(:main) }.not_to raise_error
end
end
end
end
context 'on Ci::CiDatabaseRecord' do
before do
skip_if_multiple_databases_not_setup
end
it 'raises exception' do
expect { Ci::CiDatabaseRecord.establish_connection(:ci) }.to raise_error /Cannot re-establish/
end
context 'when using with_reestablished_active_record_base' do
it 'does not raise exception' do
with_reestablished_active_record_base do
expect { Ci::CiDatabaseRecord.establish_connection(:main) }.not_to raise_error
end
end
end
end
end
context 'when trying to access connection' do
context 'when reconnect is true' do
it 'does not raise exception' do
with_reestablished_active_record_base(reconnect: true) do
expect { ActiveRecord::Base.connection.execute("SELECT 1") }.not_to raise_error # rubocop:disable Database/MultipleDatabases
end
end
end
context 'when reconnect is false' do
it 'does raise exception' do
with_reestablished_active_record_base(reconnect: false) do
expect { ActiveRecord::Base.connection.execute("SELECT 1") }.to raise_error(ActiveRecord::ConnectionNotEstablished) # rubocop:disable Database/MultipleDatabases
end
end
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