summaryrefslogtreecommitdiff
path: root/spec/lib/gitlab/database
diff options
context:
space:
mode:
Diffstat (limited to 'spec/lib/gitlab/database')
-rw-r--r--spec/lib/gitlab/database/async_indexes/index_creator_spec.rb50
-rw-r--r--spec/lib/gitlab/database/async_indexes/migration_helpers_spec.rb176
-rw-r--r--spec/lib/gitlab/database/async_indexes/postgres_async_index_spec.rb17
-rw-r--r--spec/lib/gitlab/database/async_indexes_spec.rb23
-rw-r--r--spec/lib/gitlab/database/connection_spec.rb467
-rw-r--r--spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb30
-rw-r--r--spec/lib/gitlab/database/load_balancing/host_list_spec.rb84
-rw-r--r--spec/lib/gitlab/database/load_balancing/host_spec.rb15
-rw-r--r--spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb247
-rw-r--r--spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb9
-rw-r--r--spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb67
-rw-r--r--spec/lib/gitlab/database/load_balancing/sticking_spec.rb16
-rw-r--r--spec/lib/gitlab/database/load_balancing_spec.rb150
-rw-r--r--spec/lib/gitlab/database/migration_helpers_spec.rb108
-rw-r--r--spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb97
-rw-r--r--spec/lib/gitlab/database/migrations/instrumentation_spec.rb47
-rw-r--r--spec/lib/gitlab/database/migrations/observers/query_details_spec.rb11
-rw-r--r--spec/lib/gitlab/database/migrations/observers/query_log_spec.rb12
-rw-r--r--spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb8
-rw-r--r--spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb8
-rw-r--r--spec/lib/gitlab/database/multi_threaded_migration_spec.rb43
-rw-r--r--spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb181
-rw-r--r--spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb20
-rw-r--r--spec/lib/gitlab/database/partitioning/partition_manager_spec.rb143
-rw-r--r--spec/lib/gitlab/database/partitioning/partition_monitoring_spec.rb12
-rw-r--r--spec/lib/gitlab/database/postgres_foreign_key_spec.rb41
-rw-r--r--spec/lib/gitlab/database/postgres_index_spec.rb31
-rw-r--r--spec/lib/gitlab/database/reindexing_spec.rb29
-rw-r--r--spec/lib/gitlab/database/schema_migrations/context_spec.rb61
-rw-r--r--spec/lib/gitlab/database/transaction/context_spec.rb144
-rw-r--r--spec/lib/gitlab/database/transaction/observer_spec.rb57
-rw-r--r--spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb10
-rw-r--r--spec/lib/gitlab/database/with_lock_retries_spec.rb11
33 files changed, 1859 insertions, 566 deletions
diff --git a/spec/lib/gitlab/database/async_indexes/index_creator_spec.rb b/spec/lib/gitlab/database/async_indexes/index_creator_spec.rb
new file mode 100644
index 00000000000..b4010d0fe8d
--- /dev/null
+++ b/spec/lib/gitlab/database/async_indexes/index_creator_spec.rb
@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::AsyncIndexes::IndexCreator do
+ describe '#perform' do
+ subject { described_class.new(async_index) }
+
+ let(:async_index) { create(:postgres_async_index) }
+
+ let(:index_model) { Gitlab::Database::AsyncIndexes::PostgresAsyncIndex }
+
+ let(:connection) { ApplicationRecord.connection }
+
+ context 'when the index already exists' do
+ before do
+ connection.execute(async_index.definition)
+ end
+
+ it 'skips index creation' do
+ expect(connection).not_to receive(:execute).with(/CREATE INDEX/)
+
+ subject.perform
+ end
+ end
+
+ it 'creates the index while controlling statement timeout' do
+ allow(connection).to receive(:execute).and_call_original
+ expect(connection).to receive(:execute).with("SET statement_timeout TO '32400s'").ordered.and_call_original
+ expect(connection).to receive(:execute).with(async_index.definition).ordered.and_call_original
+ expect(connection).to receive(:execute).with("RESET statement_timeout").ordered.and_call_original
+
+ subject.perform
+ end
+
+ it 'removes the index preparation record from postgres_async_indexes' do
+ expect(async_index).to receive(:destroy).and_call_original
+
+ expect { subject.perform }.to change { index_model.count }.by(-1)
+ end
+
+ it 'skips logic if not able to acquire exclusive lease' do
+ expect(subject).to receive(:try_obtain_lease).and_return(false)
+ expect(connection).not_to receive(:execute).with(/CREATE INDEX/)
+ expect(async_index).not_to receive(:destroy)
+
+ expect { subject.perform }.not_to change { index_model.count }
+ end
+ end
+end
diff --git a/spec/lib/gitlab/database/async_indexes/migration_helpers_spec.rb b/spec/lib/gitlab/database/async_indexes/migration_helpers_spec.rb
new file mode 100644
index 00000000000..ed15951dfb0
--- /dev/null
+++ b/spec/lib/gitlab/database/async_indexes/migration_helpers_spec.rb
@@ -0,0 +1,176 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::AsyncIndexes::MigrationHelpers do
+ let(:migration) { ActiveRecord::Migration.new.extend(described_class) }
+ let(:index_model) { Gitlab::Database::AsyncIndexes::PostgresAsyncIndex }
+ let(:connection) { ApplicationRecord.connection }
+ let(:table_name) { '_test_async_indexes' }
+ let(:index_name) { "index_#{table_name}_on_id" }
+
+ before do
+ allow(migration).to receive(:puts)
+ end
+
+ describe '#unprepare_async_index' do
+ let!(:async_index) { create(:postgres_async_index, name: index_name) }
+
+ context 'when the flag is enabled' do
+ before do
+ stub_feature_flags(database_async_index_creation: true)
+ end
+
+ it 'destroys the record' do
+ expect do
+ migration.unprepare_async_index(table_name, 'id')
+ end.to change { index_model.where(name: index_name).count }.by(-1)
+ end
+
+ context 'when an explicit name is given' do
+ let(:index_name) { 'my_test_async_index' }
+
+ it 'destroys the record' do
+ expect do
+ migration.unprepare_async_index(table_name, 'id', name: index_name)
+ end.to change { index_model.where(name: index_name).count }.by(-1)
+ end
+ end
+
+ context 'when the async index table does not exist' do
+ it 'does not raise an error' do
+ connection.drop_table(:postgres_async_indexes)
+
+ expect(index_model).not_to receive(:find_by)
+
+ expect { migration.unprepare_async_index(table_name, 'id') }.not_to raise_error
+ end
+ end
+ end
+
+ context 'when the feature flag is disabled' do
+ it 'does not destroy the record' do
+ stub_feature_flags(database_async_index_creation: false)
+
+ expect do
+ migration.unprepare_async_index(table_name, 'id')
+ end.not_to change { index_model.where(name: index_name).count }
+ end
+ end
+ end
+
+ describe '#unprepare_async_index_by_name' do
+ let(:index_name) { "index_#{table_name}_on_id" }
+ let!(:async_index) { create(:postgres_async_index, name: index_name) }
+
+ context 'when the flag is enabled' do
+ before do
+ stub_feature_flags(database_async_index_creation: true)
+ end
+
+ it 'destroys the record' do
+ expect do
+ migration.unprepare_async_index_by_name(table_name, index_name)
+ end.to change { index_model.where(name: index_name).count }.by(-1)
+ end
+
+ context 'when the async index table does not exist' do
+ it 'does not raise an error' do
+ connection.drop_table(:postgres_async_indexes)
+
+ expect(index_model).not_to receive(:find_by)
+
+ expect { migration.unprepare_async_index_by_name(table_name, index_name) }.not_to raise_error
+ end
+ end
+ end
+
+ context 'when the feature flag is disabled' do
+ it 'does not destroy the record' do
+ stub_feature_flags(database_async_index_creation: false)
+
+ expect do
+ migration.unprepare_async_index_by_name(table_name, index_name)
+ end.not_to change { index_model.where(name: index_name).count }
+ end
+ end
+ end
+
+ describe '#prepare_async_index' do
+ before do
+ connection.create_table(table_name)
+ end
+
+ context 'when the feature flag is enabled' do
+ before do
+ stub_feature_flags(database_async_index_creation: true)
+ end
+
+ it 'creates the record for the async index' do
+ expect do
+ migration.prepare_async_index(table_name, 'id')
+ end.to change { index_model.where(name: index_name).count }.by(1)
+
+ record = index_model.find_by(name: index_name)
+
+ expect(record.table_name).to eq(table_name)
+ expect(record.definition).to match(/CREATE INDEX CONCURRENTLY "#{index_name}"/)
+ end
+
+ context 'when an explicit name is given' do
+ let(:index_name) { 'my_async_index_name' }
+
+ it 'creates the record with the given name' do
+ expect do
+ migration.prepare_async_index(table_name, 'id', name: index_name)
+ end.to change { index_model.where(name: index_name).count }.by(1)
+
+ record = index_model.find_by(name: index_name)
+
+ expect(record.table_name).to eq(table_name)
+ expect(record.definition).to match(/CREATE INDEX CONCURRENTLY "#{index_name}"/)
+ end
+ end
+
+ context 'when the index already exists' do
+ it 'does not create the record' do
+ connection.add_index(table_name, 'id', name: index_name)
+
+ expect do
+ migration.prepare_async_index(table_name, 'id')
+ end.not_to change { index_model.where(name: index_name).count }
+ end
+ end
+
+ context 'when the record already exists' do
+ it 'does attempt to create the record' do
+ create(:postgres_async_index, table_name: table_name, name: index_name)
+
+ expect do
+ migration.prepare_async_index(table_name, 'id')
+ end.not_to change { index_model.where(name: index_name).count }
+ end
+ end
+
+ context 'when the async index table does not exist' do
+ it 'does not raise an error' do
+ connection.drop_table(:postgres_async_indexes)
+
+ expect(index_model).not_to receive(:safe_find_or_create_by!)
+
+ expect { migration.prepare_async_index(table_name, 'id') }.not_to raise_error
+ end
+ end
+ end
+
+ context 'when the feature flag is disabled' do
+ it 'does not create the record' do
+ stub_feature_flags(database_async_index_creation: false)
+
+ expect do
+ migration.prepare_async_index(table_name, 'id')
+ end.not_to change { index_model.where(name: index_name).count }
+ end
+ end
+ end
+end
diff --git a/spec/lib/gitlab/database/async_indexes/postgres_async_index_spec.rb b/spec/lib/gitlab/database/async_indexes/postgres_async_index_spec.rb
new file mode 100644
index 00000000000..434cba4edde
--- /dev/null
+++ b/spec/lib/gitlab/database/async_indexes/postgres_async_index_spec.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::AsyncIndexes::PostgresAsyncIndex, type: :model do
+ describe 'validations' do
+ let(:identifier_limit) { described_class::MAX_IDENTIFIER_LENGTH }
+ let(:definition_limit) { described_class::MAX_DEFINITION_LENGTH }
+
+ it { is_expected.to validate_presence_of(:name) }
+ it { is_expected.to validate_length_of(:name).is_at_most(identifier_limit) }
+ it { is_expected.to validate_presence_of(:table_name) }
+ it { is_expected.to validate_length_of(:table_name).is_at_most(identifier_limit) }
+ it { is_expected.to validate_presence_of(:definition) }
+ it { is_expected.to validate_length_of(:definition).is_at_most(definition_limit) }
+ end
+end
diff --git a/spec/lib/gitlab/database/async_indexes_spec.rb b/spec/lib/gitlab/database/async_indexes_spec.rb
new file mode 100644
index 00000000000..74e30ea2c4e
--- /dev/null
+++ b/spec/lib/gitlab/database/async_indexes_spec.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::AsyncIndexes do
+ describe '.create_pending_indexes!' do
+ subject { described_class.create_pending_indexes! }
+
+ before do
+ create_list(:postgres_async_index, 4)
+ end
+
+ it 'takes 2 pending indexes and creates those' do
+ Gitlab::Database::AsyncIndexes::PostgresAsyncIndex.order(:id).limit(2).each do |index|
+ creator = double('index creator')
+ expect(Gitlab::Database::AsyncIndexes::IndexCreator).to receive(:new).with(index).and_return(creator)
+ expect(creator).to receive(:perform)
+ end
+
+ subject
+ end
+ end
+end
diff --git a/spec/lib/gitlab/database/connection_spec.rb b/spec/lib/gitlab/database/connection_spec.rb
new file mode 100644
index 00000000000..5e0e6039afc
--- /dev/null
+++ b/spec/lib/gitlab/database/connection_spec.rb
@@ -0,0 +1,467 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Connection do
+ let(:connection) { described_class.new }
+
+ describe '#default_pool_size' do
+ before do
+ allow(Gitlab::Runtime).to receive(:max_threads).and_return(7)
+ end
+
+ it 'returns the max thread size plus a fixed headroom of 10' do
+ expect(connection.default_pool_size).to eq(17)
+ end
+
+ it 'returns the max thread size plus a DB_POOL_HEADROOM if this env var is present' do
+ stub_env('DB_POOL_HEADROOM', '7')
+
+ expect(connection.default_pool_size).to eq(14)
+ end
+ end
+
+ describe '#config' do
+ it 'returns a HashWithIndifferentAccess' do
+ expect(connection.config).to be_an_instance_of(HashWithIndifferentAccess)
+ end
+
+ it 'returns a default pool size' do
+ expect(connection.config).to include(pool: connection.default_pool_size)
+ end
+
+ it 'does not cache its results' do
+ a = connection.config
+ b = connection.config
+
+ expect(a).not_to equal(b)
+ end
+ end
+
+ describe '#pool_size' do
+ context 'when no explicit size is configured' do
+ it 'returns the default pool size' do
+ expect(connection).to receive(:config).and_return({ pool: nil })
+
+ expect(connection.pool_size).to eq(connection.default_pool_size)
+ end
+ end
+
+ context 'when an explicit pool size is set' do
+ it 'returns the pool size' do
+ expect(connection).to receive(:config).and_return({ pool: 4 })
+
+ expect(connection.pool_size).to eq(4)
+ end
+ end
+ end
+
+ describe '#username' do
+ context 'when a username is set' do
+ it 'returns the username' do
+ allow(connection).to receive(:config).and_return(username: 'bob')
+
+ expect(connection.username).to eq('bob')
+ end
+ end
+
+ context 'when a username is not set' do
+ it 'returns the value of the USER environment variable' do
+ allow(connection).to receive(:config).and_return(username: nil)
+ allow(ENV).to receive(:[]).with('USER').and_return('bob')
+
+ expect(connection.username).to eq('bob')
+ end
+ end
+ end
+
+ describe '#database_name' do
+ it 'returns the name of the database' do
+ allow(connection).to receive(:config).and_return(database: 'test')
+
+ expect(connection.database_name).to eq('test')
+ end
+ end
+
+ describe '#adapter_name' do
+ it 'returns the database adapter name' do
+ allow(connection).to receive(:config).and_return(adapter: 'test')
+
+ expect(connection.adapter_name).to eq('test')
+ end
+ end
+
+ describe '#human_adapter_name' do
+ context 'when the adapter is PostgreSQL' do
+ it 'returns PostgreSQL' do
+ allow(connection).to receive(:config).and_return(adapter: 'postgresql')
+
+ expect(connection.human_adapter_name).to eq('PostgreSQL')
+ end
+ end
+
+ context 'when the adapter is not PostgreSQL' do
+ it 'returns Unknown' do
+ allow(connection).to receive(:config).and_return(adapter: 'kittens')
+
+ expect(connection.human_adapter_name).to eq('Unknown')
+ end
+ end
+ end
+
+ describe '#postgresql?' do
+ context 'when using PostgreSQL' do
+ it 'returns true' do
+ allow(connection).to receive(:adapter_name).and_return('PostgreSQL')
+
+ expect(connection.postgresql?).to eq(true)
+ end
+ end
+
+ context 'when not using PostgreSQL' do
+ it 'returns false' do
+ allow(connection).to receive(:adapter_name).and_return('MySQL')
+
+ expect(connection.postgresql?).to eq(false)
+ end
+ end
+ end
+
+ describe '#db_config_with_default_pool_size' do
+ it 'returns db_config with our default pool size' do
+ allow(connection).to receive(:default_pool_size).and_return(9)
+
+ expect(connection.db_config_with_default_pool_size.pool).to eq(9)
+ end
+
+ it 'returns db_config with the correct database name' do
+ db_name = connection.scope.connection.pool.db_config.name
+
+ expect(connection.db_config_with_default_pool_size.name).to eq(db_name)
+ end
+ end
+
+ describe '#disable_prepared_statements' do
+ around do |example|
+ original_config = ::Gitlab::Database.main.config
+
+ example.run
+
+ connection.scope.establish_connection(original_config)
+ end
+
+ it 'disables prepared statements' do
+ connection.scope.establish_connection(
+ ::Gitlab::Database.main.config.merge(prepared_statements: true)
+ )
+
+ expect(connection.scope.connection.prepared_statements).to eq(true)
+
+ connection.disable_prepared_statements
+
+ expect(connection.scope.connection.prepared_statements).to eq(false)
+ end
+
+ context 'with dynamic connection pool size' do
+ before do
+ connection.scope.establish_connection(connection.config.merge(pool: 7))
+ end
+
+ it 'retains the set pool size' do
+ connection.disable_prepared_statements
+
+ expect(connection.scope.connection.prepared_statements).to eq(false)
+ expect(connection.scope.connection.pool.size).to eq(7)
+ end
+ end
+ end
+
+ describe '#db_read_only?' do
+ it 'detects a read-only database' do
+ allow(connection.scope.connection)
+ .to receive(:execute)
+ .with('SELECT pg_is_in_recovery()')
+ .and_return([{ "pg_is_in_recovery" => "t" }])
+
+ expect(connection.db_read_only?).to be_truthy
+ end
+
+ it 'detects a read-only database' do
+ allow(connection.scope.connection)
+ .to receive(:execute)
+ .with('SELECT pg_is_in_recovery()')
+ .and_return([{ "pg_is_in_recovery" => true }])
+
+ expect(connection.db_read_only?).to be_truthy
+ end
+
+ it 'detects a read-write database' do
+ allow(connection.scope.connection)
+ .to receive(:execute)
+ .with('SELECT pg_is_in_recovery()')
+ .and_return([{ "pg_is_in_recovery" => "f" }])
+
+ expect(connection.db_read_only?).to be_falsey
+ end
+
+ it 'detects a read-write database' do
+ allow(connection.scope.connection)
+ .to receive(:execute)
+ .with('SELECT pg_is_in_recovery()')
+ .and_return([{ "pg_is_in_recovery" => false }])
+
+ expect(connection.db_read_only?).to be_falsey
+ end
+ end
+
+ describe '#db_read_write?' do
+ it 'detects a read-only database' do
+ allow(connection.scope.connection)
+ .to receive(:execute)
+ .with('SELECT pg_is_in_recovery()')
+ .and_return([{ "pg_is_in_recovery" => "t" }])
+
+ expect(connection.db_read_write?).to eq(false)
+ end
+
+ it 'detects a read-only database' do
+ allow(connection.scope.connection)
+ .to receive(:execute)
+ .with('SELECT pg_is_in_recovery()')
+ .and_return([{ "pg_is_in_recovery" => true }])
+
+ expect(connection.db_read_write?).to eq(false)
+ end
+
+ it 'detects a read-write database' do
+ allow(connection.scope.connection)
+ .to receive(:execute)
+ .with('SELECT pg_is_in_recovery()')
+ .and_return([{ "pg_is_in_recovery" => "f" }])
+
+ expect(connection.db_read_write?).to eq(true)
+ end
+
+ it 'detects a read-write database' do
+ allow(connection.scope.connection)
+ .to receive(:execute)
+ .with('SELECT pg_is_in_recovery()')
+ .and_return([{ "pg_is_in_recovery" => false }])
+
+ expect(connection.db_read_write?).to eq(true)
+ end
+ end
+
+ describe '#version' do
+ around do |example|
+ connection.instance_variable_set(:@version, nil)
+ example.run
+ connection.instance_variable_set(:@version, nil)
+ end
+
+ context "on postgresql" do
+ it "extracts the version number" do
+ allow(connection)
+ .to receive(:database_version)
+ .and_return("PostgreSQL 9.4.4 on x86_64-apple-darwin14.3.0")
+
+ expect(connection.version).to eq '9.4.4'
+ end
+ end
+
+ it 'memoizes the result' do
+ count = ActiveRecord::QueryRecorder
+ .new { 2.times { connection.version } }
+ .count
+
+ expect(count).to eq(1)
+ end
+ end
+
+ describe '#postgresql_minimum_supported_version?' do
+ it 'returns false when using PostgreSQL 10' do
+ allow(connection).to receive(:version).and_return('10')
+
+ expect(connection.postgresql_minimum_supported_version?).to eq(false)
+ end
+
+ it 'returns false when using PostgreSQL 11' do
+ allow(connection).to receive(:version).and_return('11')
+
+ expect(connection.postgresql_minimum_supported_version?).to eq(false)
+ end
+
+ it 'returns true when using PostgreSQL 12' do
+ allow(connection).to receive(:version).and_return('12')
+
+ expect(connection.postgresql_minimum_supported_version?).to eq(true)
+ end
+ end
+
+ describe '#bulk_insert' do
+ before do
+ allow(connection).to receive(:connection).and_return(dummy_connection)
+ allow(dummy_connection).to receive(:quote_column_name, &:itself)
+ allow(dummy_connection).to receive(:quote, &:itself)
+ allow(dummy_connection).to receive(:execute)
+ end
+
+ let(:dummy_connection) { double(:connection) }
+
+ let(:rows) do
+ [
+ { a: 1, b: 2, c: 3 },
+ { c: 6, a: 4, b: 5 }
+ ]
+ end
+
+ it 'does nothing with empty rows' do
+ expect(dummy_connection).not_to receive(:execute)
+
+ connection.bulk_insert('test', [])
+ end
+
+ it 'uses the ordering from the first row' do
+ expect(dummy_connection).to receive(:execute) do |sql|
+ expect(sql).to include('(1, 2, 3)')
+ expect(sql).to include('(4, 5, 6)')
+ end
+
+ connection.bulk_insert('test', rows)
+ end
+
+ it 'quotes column names' do
+ expect(dummy_connection).to receive(:quote_column_name).with(:a)
+ expect(dummy_connection).to receive(:quote_column_name).with(:b)
+ expect(dummy_connection).to receive(:quote_column_name).with(:c)
+
+ connection.bulk_insert('test', rows)
+ end
+
+ it 'quotes values' do
+ 1.upto(6) do |i|
+ expect(dummy_connection).to receive(:quote).with(i)
+ end
+
+ connection.bulk_insert('test', rows)
+ end
+
+ it 'does not quote values of a column in the disable_quote option' do
+ [1, 2, 4, 5].each do |i|
+ expect(dummy_connection).to receive(:quote).with(i)
+ end
+
+ connection.bulk_insert('test', rows, disable_quote: :c)
+ end
+
+ it 'does not quote values of columns in the disable_quote option' do
+ [2, 5].each do |i|
+ expect(dummy_connection).to receive(:quote).with(i)
+ end
+
+ connection.bulk_insert('test', rows, disable_quote: [:a, :c])
+ end
+
+ it 'handles non-UTF-8 data' do
+ expect { connection.bulk_insert('test', [{ a: "\255" }]) }.not_to raise_error
+ end
+
+ context 'when using PostgreSQL' do
+ it 'allows the returning of the IDs of the inserted rows' do
+ result = double(:result, values: [['10']])
+
+ expect(dummy_connection)
+ .to receive(:execute)
+ .with(/RETURNING id/)
+ .and_return(result)
+
+ ids = connection
+ .bulk_insert('test', [{ number: 10 }], return_ids: true)
+
+ expect(ids).to eq([10])
+ end
+
+ it 'allows setting the upsert to do nothing' do
+ expect(dummy_connection)
+ .to receive(:execute)
+ .with(/ON CONFLICT DO NOTHING/)
+
+ connection
+ .bulk_insert('test', [{ number: 10 }], on_conflict: :do_nothing)
+ end
+ end
+ end
+
+ describe '#cached_column_exists?' do
+ it 'only retrieves data once' do
+ expect(connection.scope.connection)
+ .to receive(:columns)
+ .once.and_call_original
+
+ 2.times do
+ expect(connection.cached_column_exists?(:projects, :id)).to be_truthy
+ expect(connection.cached_column_exists?(:projects, :bogus_column)).to be_falsey
+ end
+ end
+ end
+
+ describe '#cached_table_exists?' do
+ it 'only retrieves data once per table' do
+ expect(connection.scope.connection)
+ .to receive(:data_source_exists?)
+ .with(:projects)
+ .once.and_call_original
+
+ expect(connection.scope.connection)
+ .to receive(:data_source_exists?)
+ .with(:bogus_table_name)
+ .once.and_call_original
+
+ 2.times do
+ expect(connection.cached_table_exists?(:projects)).to be_truthy
+ expect(connection.cached_table_exists?(:bogus_table_name)).to be_falsey
+ end
+ end
+
+ it 'returns false when database does not exist' do
+ expect(connection.scope).to receive(:connection) do
+ raise ActiveRecord::NoDatabaseError, 'broken'
+ end
+
+ expect(connection.cached_table_exists?(:projects)).to be(false)
+ end
+ end
+
+ describe '#exists?' do
+ it 'returns true if `ActiveRecord::Base.connection` succeeds' do
+ expect(connection.scope).to receive(:connection)
+
+ expect(connection.exists?).to be(true)
+ end
+
+ it 'returns false if `ActiveRecord::Base.connection` fails' do
+ expect(connection.scope).to receive(:connection) do
+ raise ActiveRecord::NoDatabaseError, 'broken'
+ end
+
+ expect(connection.exists?).to be(false)
+ end
+ end
+
+ describe '#system_id' do
+ it 'returns the PostgreSQL system identifier' do
+ expect(connection.system_id).to be_an_instance_of(Integer)
+ end
+ end
+
+ describe '#get_write_location' do
+ it 'returns a string' do
+ expect(connection.get_write_location(connection.scope.connection))
+ .to be_a(String)
+ end
+
+ it 'returns nil if there are no results' do
+ expect(connection.get_write_location(double(select_all: []))).to be_nil
+ end
+ end
+end
diff --git a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
index 015dd2ba8d2..0ca99ec9acf 100644
--- a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
describe '#select' do
it 'performs a read' do
- expect(proxy).to receive(:read_using_load_balancer).with(:select, ['foo'])
+ expect(proxy).to receive(:read_using_load_balancer).with(:select, 'foo')
proxy.select('foo')
end
@@ -26,7 +26,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
arel = double(:arel)
expect(proxy).to receive(:read_using_load_balancer)
- .with(:select_all, [arel, 'foo', []])
+ .with(:select_all, arel, 'foo', [])
proxy.select_all(arel, 'foo')
end
@@ -37,7 +37,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
arel = double(:arel, locked: true)
expect(proxy).to receive(:write_using_load_balancer)
- .with(:select_all, [arel, 'foo', []], sticky: true)
+ .with(:select_all, arel, 'foo', [], sticky: true)
proxy.select_all(arel, 'foo')
end
@@ -48,7 +48,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
describe "#{name}" do
it 'runs the query on the replica' do
expect(proxy).to receive(:read_using_load_balancer)
- .with(name, ['foo'])
+ .with(name, 'foo')
proxy.send(name, 'foo')
end
@@ -59,7 +59,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
describe "#{name}" do
it 'runs the query on the primary and sticks to it' do
expect(proxy).to receive(:write_using_load_balancer)
- .with(name, ['foo'], sticky: true)
+ .with(name, 'foo', sticky: true)
proxy.send(name, 'foo')
end
@@ -187,7 +187,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
describe '#method_missing' do
it 'runs the query on the primary without sticking to it' do
expect(proxy).to receive(:write_using_load_balancer)
- .with(:foo, ['foo'])
+ .with(:foo, 'foo')
proxy.foo('foo')
end
@@ -197,7 +197,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
expect(proxy).to receive(:write_using_load_balancer).and_call_original
- expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) }
+ expect { proxy.case_sensitive_comparison(:table, :attribute, :column, value: :value, format: :format) }
.not_to raise_error
end
@@ -212,7 +212,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
end
it 'runs the query on the replica' do
- expect(proxy).to receive(:read_using_load_balancer).with(:foo, ['foo'])
+ expect(proxy).to receive(:read_using_load_balancer).with(:foo, 'foo')
proxy.foo('foo')
end
@@ -222,7 +222,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
expect(proxy).to receive(:read_using_load_balancer).and_call_original
- expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) }
+ expect { proxy.case_sensitive_comparison(:table, :attribute, :column, value: :value, format: :format) }
.not_to raise_error
end
end
@@ -245,7 +245,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
expect(connection).to receive(:foo).with('foo')
expect(proxy.load_balancer).to receive(:read).and_yield(connection)
- proxy.read_using_load_balancer(:foo, ['foo'])
+ proxy.read_using_load_balancer(:foo, 'foo')
end
end
@@ -257,7 +257,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
expect(connection).to receive(:foo).with('foo')
expect(proxy.load_balancer).to receive(:read).and_yield(connection)
- proxy.read_using_load_balancer(:foo, ['foo'])
+ proxy.read_using_load_balancer(:foo, 'foo')
end
end
@@ -269,7 +269,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
expect(connection).to receive(:foo).with('foo')
expect(proxy.load_balancer).to receive(:read).and_yield(connection)
- proxy.read_using_load_balancer(:foo, ['foo'])
+ proxy.read_using_load_balancer(:foo, 'foo')
end
end
@@ -283,7 +283,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
expect(proxy.load_balancer).to receive(:read_write)
.and_yield(connection)
- proxy.read_using_load_balancer(:foo, ['foo'])
+ proxy.read_using_load_balancer(:foo, 'foo')
end
end
end
@@ -302,7 +302,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
expect(connection).to receive(:foo).with('foo')
expect(session).not_to receive(:write!)
- proxy.write_using_load_balancer(:foo, ['foo'])
+ proxy.write_using_load_balancer(:foo, 'foo')
end
it 'sticks to the primary when sticking is enabled' do
@@ -310,7 +310,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
expect(connection).to receive(:foo).with('foo')
expect(session).to receive(:write!)
- proxy.write_using_load_balancer(:foo, ['foo'], sticky: true)
+ proxy.write_using_load_balancer(:foo, 'foo', sticky: true)
end
end
end
diff --git a/spec/lib/gitlab/database/load_balancing/host_list_spec.rb b/spec/lib/gitlab/database/load_balancing/host_list_spec.rb
index 873b599f84d..ad4ca18d5e6 100644
--- a/spec/lib/gitlab/database/load_balancing/host_list_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/host_list_spec.rb
@@ -3,25 +3,17 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing::HostList do
- def expect_metrics(hosts)
- expect(Gitlab::Metrics.registry.get(:db_load_balancing_hosts).get({})).to eq(hosts)
- end
-
- before do
- allow(Gitlab::Database)
- .to receive(:create_connection_pool)
- .and_return(ActiveRecord::Base.connection_pool)
- end
-
+ let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
let(:load_balancer) { double(:load_balancer) }
let(:host_count) { 2 }
+ let(:hosts) { Array.new(host_count) { Gitlab::Database::LoadBalancing::Host.new(db_host, load_balancer, port: 5432) } }
+ let(:host_list) { described_class.new(hosts) }
- let(:host_list) do
- hosts = Array.new(host_count) do
- Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer, port: 5432)
+ before do
+ # each call generate a new replica pool
+ allow(load_balancer).to receive(:create_replica_connection_pool) do
+ double(:replica_connection_pool)
end
-
- described_class.new(hosts)
end
describe '#initialize' do
@@ -42,8 +34,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
context 'with ports' do
it 'returns the host names of all hosts' do
hosts = [
- ['localhost', 5432],
- ['localhost', 5432]
+ [db_host, 5432],
+ [db_host, 5432]
]
expect(host_list.host_names_and_ports).to eq(hosts)
@@ -51,18 +43,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
end
context 'without ports' do
- let(:host_list) do
- hosts = Array.new(2) do
- Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer)
- end
-
- described_class.new(hosts)
- end
+ let(:hosts) { Array.new(2) { Gitlab::Database::LoadBalancing::Host.new(db_host, load_balancer) } }
it 'returns the host names of all hosts' do
hosts = [
- ['localhost', nil],
- ['localhost', nil]
+ [db_host, nil],
+ [db_host, nil]
]
expect(host_list.host_names_and_ports).to eq(hosts)
@@ -70,48 +56,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
end
end
- describe '#manage_pool?' do
- before do
- allow(Gitlab::Database).to receive(:create_connection_pool) { double(:connection) }
- end
-
- context 'when the testing pool belongs to one host of the host list' do
- it 'returns true' do
- pool = host_list.hosts.first.pool
-
- expect(host_list.manage_pool?(pool)).to be(true)
- end
- end
-
- context 'when the testing pool belongs to a former host of the host list' do
- it 'returns false' do
- pool = host_list.hosts.first.pool
- host_list.hosts = [
- Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
- ]
-
- expect(host_list.manage_pool?(pool)).to be(false)
- end
- end
-
- context 'when the testing pool belongs to a new host of the host list' do
- it 'returns true' do
- host = Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
- host_list.hosts = [host]
-
- expect(host_list.manage_pool?(host.pool)).to be(true)
- end
- end
-
- context 'when the testing pool does not have any relation with the host list' do
- it 'returns false' do
- host = Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
-
- expect(host_list.manage_pool?(host.pool)).to be(false)
- end
- end
- end
-
describe '#hosts' do
it 'returns a copy of the host' do
first = host_list.hosts
@@ -185,4 +129,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
end
end
end
+
+ def expect_metrics(hosts)
+ expect(Gitlab::Metrics.registry.get(:db_load_balancing_hosts).get({})).to eq(hosts)
+ end
end
diff --git a/spec/lib/gitlab/database/load_balancing/host_spec.rb b/spec/lib/gitlab/database/load_balancing/host_spec.rb
index 4dfddef68c8..f42ac8be1bb 100644
--- a/spec/lib/gitlab/database/load_balancing/host_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/host_spec.rb
@@ -3,15 +3,16 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing::Host do
- let(:load_balancer) do
- Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[localhost])
- end
+ let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new }
- let(:host) { load_balancer.host_list.hosts.first }
+ let(:host) do
+ Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer)
+ end
before do
- allow(Gitlab::Database).to receive(:create_connection_pool)
- .and_return(ActiveRecord::Base.connection_pool)
+ allow(load_balancer).to receive(:create_replica_connection_pool) do
+ ActiveRecord::Base.connection_pool
+ end
end
def raise_and_wrap(wrapper, original)
@@ -63,7 +64,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Host do
expect(host.pool)
.to receive(:disconnect!)
- host.disconnect!(1)
+ host.disconnect!(timeout: 1)
end
end
diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
index b82b8d9a311..c647f5a8f5d 100644
--- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -3,20 +3,22 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
- let(:pool) { Gitlab::Database.create_connection_pool(2) }
let(:conflict_error) { Class.new(RuntimeError) }
-
- let(:lb) { described_class.new(%w(localhost localhost)) }
+ let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
+ let(:lb) { described_class.new([db_host, db_host]) }
+ let(:request_cache) { lb.send(:request_cache) }
before do
- allow(Gitlab::Database).to receive(:create_connection_pool)
- .and_return(pool)
stub_const(
'Gitlab::Database::LoadBalancing::LoadBalancer::PG::TRSerializationFailure',
conflict_error
)
end
+ after do |example|
+ lb.disconnect!(timeout: 0) unless example.metadata[:skip_disconnect]
+ end
+
def raise_and_wrap(wrapper, original)
raise original
rescue original.class
@@ -123,8 +125,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
describe '#read_write' do
it 'yields a connection for a write' do
- expect { |b| lb.read_write(&b) }
- .to yield_with_args(ActiveRecord::Base.retrieve_connection)
+ connection = ActiveRecord::Base.connection_pool.connection
+
+ expect { |b| lb.read_write(&b) }.to yield_with_args(connection)
end
it 'uses a retry with exponential backoffs' do
@@ -134,140 +137,30 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
end
end
- describe '#db_role_for_connection' do
- context 'when the load balancer creates the connection with #read' do
- it 'returns :replica' do
- role = nil
- lb.read do |connection|
- role = lb.db_role_for_connection(connection)
- end
-
- expect(role).to be(:replica)
- end
- end
-
- context 'when the load balancer uses nested #read' do
- it 'returns :replica' do
- roles = []
- lb.read do |connection_1|
- lb.read do |connection_2|
- roles << lb.db_role_for_connection(connection_2)
- end
- roles << lb.db_role_for_connection(connection_1)
- end
-
- expect(roles).to eq([:replica, :replica])
- end
- end
-
- context 'when the load balancer creates the connection with #read_write' do
- it 'returns :primary' do
- role = nil
- lb.read_write do |connection|
- role = lb.db_role_for_connection(connection)
- end
-
- expect(role).to be(:primary)
- end
- end
-
- context 'when the load balancer uses nested #read_write' do
- it 'returns :primary' do
- roles = []
- lb.read_write do |connection_1|
- lb.read_write do |connection_2|
- roles << lb.db_role_for_connection(connection_2)
- end
- roles << lb.db_role_for_connection(connection_1)
- end
-
- expect(roles).to eq([:primary, :primary])
- end
- end
-
- context 'when the load balancer falls back the connection creation to primary' do
- it 'returns :primary' do
- allow(lb).to receive(:serialization_failure?).and_return(true)
-
- role = nil
- raised = 7 # 2 hosts = 6 retries
-
- lb.read do |connection|
- if raised > 0
- raised -= 1
- raise
- end
-
- role = lb.db_role_for_connection(connection)
- end
-
- expect(role).to be(:primary)
- end
- end
-
- context 'when the load balancer uses replica after recovery from a failure' do
- it 'returns :replica' do
- allow(lb).to receive(:connection_error?).and_return(true)
-
- role = nil
- raised = false
-
- lb.read do |connection|
- unless raised
- raised = true
- raise
- end
-
- role = lb.db_role_for_connection(connection)
- end
-
- expect(role).to be(:replica)
- end
- end
-
- context 'when the connection comes from a pool managed by the host list' do
- it 'returns :replica' do
- connection = double(:connection)
- allow(connection).to receive(:pool).and_return(lb.host_list.hosts.first.pool)
-
- expect(lb.db_role_for_connection(connection)).to be(:replica)
- end
- end
-
- context 'when the connection comes from the primary pool' do
- it 'returns :primary' do
- connection = double(:connection)
- allow(connection).to receive(:pool).and_return(ActiveRecord::Base.connection_pool)
-
- expect(lb.db_role_for_connection(connection)).to be(:primary)
- end
- end
-
- context 'when the connection does not come from any known pool' do
- it 'returns nil' do
- connection = double(:connection)
- pool = double(:connection_pool)
- allow(connection).to receive(:pool).and_return(pool)
-
- expect(lb.db_role_for_connection(connection)).to be(nil)
- end
- end
- end
-
describe '#host' do
it 'returns the secondary host to use' do
expect(lb.host).to be_an_instance_of(Gitlab::Database::LoadBalancing::Host)
end
it 'stores the host in a thread-local variable' do
- RequestStore.delete(described_class::CACHE_KEY)
- RequestStore.delete(described_class::VALID_HOSTS_CACHE_KEY)
+ request_cache.delete(described_class::CACHE_KEY)
expect(lb.host_list).to receive(:next).once.and_call_original
lb.host
lb.host
end
+
+ it 'does not create conflicts with other load balancers when caching hosts' do
+ lb1 = described_class.new([db_host, db_host], ActiveRecord::Base)
+ lb2 = described_class.new([db_host, db_host], Ci::CiDatabaseRecord)
+
+ host1 = lb1.host
+ host2 = lb2.host
+
+ expect(lb1.send(:request_cache)[described_class::CACHE_KEY]).to eq(host1)
+ expect(lb2.send(:request_cache)[described_class::CACHE_KEY]).to eq(host2)
+ end
end
describe '#release_host' do
@@ -278,8 +171,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
lb.release_host
- expect(RequestStore[described_class::CACHE_KEY]).to be_nil
- expect(RequestStore[described_class::VALID_HOSTS_CACHE_KEY]).to be_nil
+ expect(request_cache[described_class::CACHE_KEY]).to be_nil
end
end
@@ -414,89 +306,76 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
end
end
- describe '#select_caught_up_hosts' do
+ describe '#select_up_to_date_host' do
let(:location) { 'AB/12345'}
let(:hosts) { lb.host_list.hosts }
- let(:valid_host_list) { RequestStore[described_class::VALID_HOSTS_CACHE_KEY] }
- let(:valid_hosts) { valid_host_list.hosts }
+ let(:set_host) { request_cache[described_class::CACHE_KEY] }
- subject { lb.select_caught_up_hosts(location) }
-
- context 'when all replicas are caught up' do
- before do
- expect(hosts).to all(receive(:caught_up?).with(location).and_return(true))
- end
-
- it 'returns true and sets all hosts to valid' do
- expect(subject).to be true
- expect(valid_host_list).to be_a(Gitlab::Database::LoadBalancing::HostList)
- expect(valid_hosts).to contain_exactly(*hosts)
- end
- end
+ subject { lb.select_up_to_date_host(location) }
context 'when none of the replicas are caught up' do
before do
expect(hosts).to all(receive(:caught_up?).with(location).and_return(false))
end
- it 'returns false and does not set the valid hosts' do
+ it 'returns false and does not update the host thread-local variable' do
expect(subject).to be false
- expect(valid_host_list).to be_nil
+ expect(set_host).to be_nil
end
end
- context 'when one of the replicas is caught up' do
+ context 'when any of the replicas is caught up' do
before do
- expect(hosts[0]).to receive(:caught_up?).with(location).and_return(false)
+ # `allow` for non-caught up host, because we may not even check it, if will find the caught up one earlier
+ allow(hosts[0]).to receive(:caught_up?).with(location).and_return(false)
expect(hosts[1]).to receive(:caught_up?).with(location).and_return(true)
end
- it 'returns true and sets one host to valid' do
+ it 'returns true and sets host thread-local variable' do
expect(subject).to be true
- expect(valid_host_list).to be_a(Gitlab::Database::LoadBalancing::HostList)
- expect(valid_hosts).to contain_exactly(hosts[1])
- end
-
- it 'host always returns the caught-up replica' do
- subject
-
- 3.times do
- expect(lb.host).to eq(hosts[1])
- RequestStore.delete(described_class::CACHE_KEY)
- end
+ expect(set_host).to eq(hosts[1])
end
end
end
- describe '#select_up_to_date_host' do
- let(:location) { 'AB/12345'}
- let(:hosts) { lb.host_list.hosts }
- let(:set_host) { RequestStore[described_class::CACHE_KEY] }
+ describe '#create_replica_connection_pool' do
+ it 'creates a new connection pool with specific pool size and name' do
+ with_replica_pool(5, 'other_host') do |replica_pool|
+ expect(replica_pool)
+ .to be_kind_of(ActiveRecord::ConnectionAdapters::ConnectionPool)
- subject { lb.select_up_to_date_host(location) }
-
- context 'when none of the replicas are caught up' do
- before do
- expect(hosts).to all(receive(:caught_up?).with(location).and_return(false))
+ expect(replica_pool.db_config.host).to eq('other_host')
+ expect(replica_pool.db_config.pool).to eq(5)
+ expect(replica_pool.db_config.name).to end_with("_replica")
end
+ end
- it 'returns false and does not update the host thread-local variable' do
- expect(subject).to be false
- expect(set_host).to be_nil
+ it 'allows setting of a custom hostname and port' do
+ with_replica_pool(5, 'other_host', 5432) do |replica_pool|
+ expect(replica_pool.db_config.host).to eq('other_host')
+ expect(replica_pool.db_config.configuration_hash[:port]).to eq(5432)
end
end
- context 'when any of the replicas is caught up' do
- before do
- # `allow` for non-caught up host, because we may not even check it, if will find the caught up one earlier
- allow(hosts[0]).to receive(:caught_up?).with(location).and_return(false)
- expect(hosts[1]).to receive(:caught_up?).with(location).and_return(true)
- end
+ it 'does not modify connection class pool' do
+ expect { with_replica_pool(5) { } }.not_to change { ActiveRecord::Base.connection_pool }
+ end
- it 'returns true and sets host thread-local variable' do
- expect(subject).to be true
- expect(set_host).to eq(hosts[1])
+ def with_replica_pool(*args)
+ pool = lb.create_replica_connection_pool(*args)
+ yield pool
+ ensure
+ pool&.disconnect!
+ end
+ end
+
+ describe '#disconnect!' do
+ it 'calls disconnect on all hosts with a timeout', :skip_disconnect do
+ expect_next_instances_of(Gitlab::Database::LoadBalancing::Host, 2) do |host|
+ expect(host).to receive(:disconnect!).with(timeout: 30)
end
+
+ lb.disconnect!(timeout: 30)
end
end
end
diff --git a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
index 9381ffa59fe..ea0c7f781fd 100644
--- a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
@@ -183,18 +183,17 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
describe '#clear' do
it 'clears the currently used host and session' do
lb = double(:lb)
- session = double(:session)
+ session = spy(:session)
allow(middleware).to receive(:load_balancer).and_return(lb)
expect(lb).to receive(:release_host)
- stub_const('Gitlab::Database::LoadBalancing::RackMiddleware::Session',
- session)
-
- expect(session).to receive(:clear_session)
+ stub_const('Gitlab::Database::LoadBalancing::Session', session)
middleware.clear
+
+ expect(session).to have_received(:clear_session)
end
end
diff --git a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
index 7fc7b5e8d11..a27341a3324 100644
--- a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
@@ -3,8 +3,14 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
+ let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new([]) }
let(:service) do
- described_class.new(nameserver: 'localhost', port: 8600, record: 'foo')
+ described_class.new(
+ nameserver: 'localhost',
+ port: 8600,
+ record: 'foo',
+ load_balancer: load_balancer
+ )
end
before do
@@ -18,7 +24,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
describe '#initialize' do
describe ':record_type' do
- subject { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) }
+ subject do
+ described_class.new(
+ nameserver: 'localhost',
+ port: 8600,
+ record: 'foo',
+ record_type: record_type,
+ load_balancer: load_balancer
+ )
+ end
context 'with a supported type' do
let(:record_type) { 'SRV' }
@@ -44,21 +58,17 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
end
it 'starts service discovery in a new thread' do
- expect(service)
- .to receive(:refresh_if_necessary)
- .and_return(5)
-
- expect(service)
- .to receive(:rand)
- .and_return(2)
+ expect(Thread).to receive(:new).ordered.and_call_original # Thread starts
- expect(service)
- .to receive(:sleep)
- .with(7)
+ expect(service).to receive(:perform_service_discovery).ordered.and_return(5)
+ expect(service).to receive(:rand).ordered.and_return(2)
+ expect(service).to receive(:sleep).ordered.with(7) # Sleep runs after thread starts
service.start.join
end
+ end
+ describe '#perform_service_discovery' do
it 'reports exceptions to Sentry' do
error = StandardError.new
@@ -70,15 +80,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
.to receive(:track_exception)
.with(error)
- expect(service)
- .to receive(:rand)
- .and_return(2)
-
- expect(service)
- .to receive(:sleep)
- .with(62)
-
- service.start.join
+ service.perform_service_discovery
end
end
@@ -155,14 +157,23 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
expect(host)
.to receive(:disconnect!)
- .with(2)
+ .with(timeout: 2)
service.replace_hosts([address_bar])
end
end
describe '#addresses_from_dns' do
- let(:service) { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) }
+ let(:service) do
+ described_class.new(
+ nameserver: 'localhost',
+ port: 8600,
+ record: 'foo',
+ record_type: record_type,
+ load_balancer: load_balancer
+ )
+ end
+
let(:packet) { double(:packet, answer: [res1, res2]) }
before do
@@ -234,13 +245,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
end
describe '#addresses_from_load_balancer' do
- it 'returns the ordered host names of the load balancer' do
- load_balancer = Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[b a])
-
- allow(service)
- .to receive(:load_balancer)
- .and_return(load_balancer)
+ let(:load_balancer) do
+ Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[b a])
+ end
+ it 'returns the ordered host names of the load balancer' do
addresses = [
described_class::Address.new('a'),
described_class::Address.new('b')
diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
index 53445d73756..cf52e59db3a 100644
--- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
@@ -237,7 +237,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
context 'when write location is nil' do
before do
- allow(Gitlab::Database).to receive(:get_write_location).and_return(nil)
+ allow(Gitlab::Database.main).to receive(:get_write_location).and_return(nil)
end
it 'does not update the write location' do
@@ -313,7 +313,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
end
it 'returns false and does not try to find caught up hosts' do
- expect(described_class).not_to receive(:select_caught_up_hosts)
+ expect(lb).not_to receive(:select_up_to_date_host)
expect(described_class.select_caught_up_replicas(:project, 42)).to be false
end
end
@@ -329,18 +329,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
expect(described_class).to receive(:unstick).with(:project, 42)
expect(described_class.select_caught_up_replicas(:project, 42)).to be true
end
-
- context 'when :load_balancing_refine_load_balancer_methods FF is disabled' do
- before do
- stub_feature_flags(load_balancing_refine_load_balancer_methods: false)
- end
-
- it 'returns true, selects hosts, and unsticks if any secondary has caught up' do
- expect(lb).to receive(:select_caught_up_hosts).and_return(true)
- expect(described_class).to receive(:unstick).with(:project, 42)
- expect(described_class.select_caught_up_replicas(:project, 42)).to be true
- end
- end
end
end
end
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
index 94717a10492..6ec8e0516f6 100644
--- a/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -3,25 +3,28 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing do
- include_context 'clear DB Load Balancing configuration'
+ describe '.proxy' do
+ before do
+ @previous_proxy = ActiveRecord::Base.load_balancing_proxy
- before do
- stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'true')
- end
+ ActiveRecord::Base.load_balancing_proxy = connection_proxy
+ end
+
+ after do
+ ActiveRecord::Base.load_balancing_proxy = @previous_proxy
+ end
- describe '.proxy' do
context 'when configured' do
- before do
- allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
- subject.configure_proxy
- end
+ let(:connection_proxy) { double(:connection_proxy) }
it 'returns the connection proxy' do
- expect(subject.proxy).to be_an_instance_of(subject::ConnectionProxy)
+ expect(subject.proxy).to eq(connection_proxy)
end
end
context 'when not configured' do
+ let(:connection_proxy) { nil }
+
it 'returns nil' do
expect(subject.proxy).to be_nil
end
@@ -40,9 +43,9 @@ RSpec.describe Gitlab::Database::LoadBalancing do
it 'returns a Hash' do
lb_config = { 'hosts' => %w(foo) }
- original_db_config = Gitlab::Database.config
+ original_db_config = Gitlab::Database.main.config
modified_db_config = original_db_config.merge(load_balancing: lb_config)
- expect(Gitlab::Database).to receive(:config).and_return(modified_db_config)
+ expect(Gitlab::Database.main).to receive(:config).and_return(modified_db_config)
expect(described_class.configuration).to eq(lb_config)
end
@@ -132,7 +135,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
describe '.enable?' do
before do
- clear_load_balancing_configuration
allow(described_class).to receive(:hosts).and_return(%w(foo))
end
@@ -173,10 +175,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
describe '.configured?' do
- before do
- clear_load_balancing_configuration
- end
-
it 'returns true when Sidekiq is being used' do
allow(described_class).to receive(:hosts).and_return(%w(foo))
allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
@@ -207,12 +205,27 @@ RSpec.describe Gitlab::Database::LoadBalancing do
describe '.configure_proxy' do
it 'configures the connection proxy' do
- allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+ allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
described_class.configure_proxy
- expect(ActiveRecord::Base.singleton_class).to have_received(:prepend)
- .with(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
+ expect(ActiveRecord::Base).to have_received(:load_balancing_proxy=)
+ .with(Gitlab::Database::LoadBalancing::ConnectionProxy)
+ end
+
+ context 'when service discovery is enabled' do
+ let(:service_discovery) { double(Gitlab::Database::LoadBalancing::ServiceDiscovery) }
+
+ it 'runs initial service discovery when configuring the connection proxy' do
+ allow(described_class)
+ .to receive(:configuration)
+ .and_return('discover' => { 'record' => 'foo' })
+
+ expect(Gitlab::Database::LoadBalancing::ServiceDiscovery).to receive(:new).and_return(service_discovery)
+ expect(service_discovery).to receive(:perform_service_discovery)
+
+ described_class.configure_proxy
+ end
end
end
@@ -298,59 +311,46 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
describe '.db_role_for_connection' do
- let(:connection) { double(:conneciton) }
-
context 'when the load balancing is not configured' do
- before do
- allow(described_class).to receive(:enable?).and_return(false)
- end
+ let(:connection) { ActiveRecord::Base.connection }
it 'returns primary' do
- expect(described_class.db_role_for_connection(connection)).to be(:primary)
+ expect(described_class.db_role_for_connection(connection)).to eq(:primary)
end
end
- context 'when the load balancing is configured' do
- let(:proxy) { described_class::ConnectionProxy.new(%w(foo)) }
- let(:load_balancer) { described_class::LoadBalancer.new(%w(foo)) }
-
- before do
- allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+ context 'when the NullPool is used for connection' do
+ let(:pool) { ActiveRecord::ConnectionAdapters::NullPool.new }
+ let(:connection) { double(:connection, pool: pool) }
- allow(described_class).to receive(:enable?).and_return(true)
- allow(described_class).to receive(:proxy).and_return(proxy)
- allow(proxy).to receive(:load_balancer).and_return(load_balancer)
-
- subject.configure_proxy(proxy)
+ it 'returns unknown' do
+ expect(described_class.db_role_for_connection(connection)).to eq(:unknown)
end
+ end
- context 'when the load balancer returns :replica' do
- it 'returns :replica' do
- allow(load_balancer).to receive(:db_role_for_connection).and_return(:replica)
-
- expect(described_class.db_role_for_connection(connection)).to be(:replica)
+ context 'when the load balancing is configured' do
+ let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
+ let(:proxy) { described_class::ConnectionProxy.new([db_host]) }
- expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+ context 'when a proxy connection is used' do
+ it 'returns :unknown' do
+ expect(described_class.db_role_for_connection(proxy)).to eq(:unknown)
end
end
- context 'when the load balancer returns :primary' do
- it 'returns :primary' do
- allow(load_balancer).to receive(:db_role_for_connection).and_return(:primary)
-
- expect(described_class.db_role_for_connection(connection)).to be(:primary)
-
- expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+ context 'when a read connection is used' do
+ it 'returns :replica' do
+ proxy.load_balancer.read do |connection|
+ expect(described_class.db_role_for_connection(connection)).to eq(:replica)
+ end
end
end
- context 'when the load balancer returns nil' do
- it 'returns nil' do
- allow(load_balancer).to receive(:db_role_for_connection).and_return(nil)
-
- expect(described_class.db_role_for_connection(connection)).to be(nil)
-
- expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+ context 'when a read_write connection is used' do
+ it 'returns :primary' do
+ proxy.load_balancer.read_write do |connection|
+ expect(described_class.db_role_for_connection(connection)).to eq(:primary)
+ end
end
end
end
@@ -366,7 +366,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
# - In each test, we listen to the SQL queries (via sql.active_record
# instrumentation) while triggering real queries from the defined model.
# - We assert the desinations (replica/primary) of the queries in order.
- describe 'LoadBalancing integration tests', :delete do
+ describe 'LoadBalancing integration tests', :db_load_balancing, :delete do
before(:all) do
ActiveRecord::Schema.define do
create_table :load_balancing_test, force: true do |t|
@@ -381,30 +381,14 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
end
- shared_context 'LoadBalancing setup' do
- let(:development_db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'development').first.configuration_hash }
- let(:hosts) { [development_db_config[:host]] }
- let(:model) do
- Class.new(ApplicationRecord) do
- self.table_name = "load_balancing_test"
- end
+ let(:model) do
+ Class.new(ApplicationRecord) do
+ self.table_name = "load_balancing_test"
end
+ end
- before do
- # Preloading testing class
- model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
-
- # Setup load balancing
- clear_load_balancing_configuration
- allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
- subject.configure_proxy(::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts))
-
- original_db_config = Gitlab::Database.config
- modified_db_config = original_db_config.merge(load_balancing: { hosts: hosts })
- allow(Gitlab::Database).to receive(:config).and_return(modified_db_config)
-
- ::Gitlab::Database::LoadBalancing::Session.clear_session
- end
+ before do
+ model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
end
where(:queries, :include_transaction, :expected_results) do
@@ -715,8 +699,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
with_them do
- include_context 'LoadBalancing setup'
-
it 'redirects queries to the right roles' do
roles = []
@@ -785,8 +767,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
with_them do
- include_context 'LoadBalancing setup'
-
it 'redirects queries to the right roles' do
roles = []
@@ -805,8 +785,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
context 'a write inside a transaction inside fallback_to_replicas_for_ambiguous_queries block' do
- include_context 'LoadBalancing setup'
-
it 'raises an exception' do
expect do
::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 8e25f9249fe..9f9aef77de7 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -278,6 +278,16 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
model.add_concurrent_index(:users, :foo, unique: true)
end
+
+ it 'unprepares the async index creation' do
+ expect(model).to receive(:add_index)
+ .with(:users, :foo, algorithm: :concurrently)
+
+ expect(model).to receive(:unprepare_async_index)
+ .with(:users, :foo, algorithm: :concurrently)
+
+ model.add_concurrent_index(:users, :foo)
+ end
end
context 'inside a transaction' do
@@ -314,6 +324,16 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
model.remove_concurrent_index(:users, :foo, unique: true)
end
+ it 'unprepares the async index creation' do
+ expect(model).to receive(:remove_index)
+ .with(:users, { algorithm: :concurrently, column: :foo })
+
+ expect(model).to receive(:unprepare_async_index)
+ .with(:users, :foo, { algorithm: :concurrently })
+
+ model.remove_concurrent_index(:users, :foo)
+ end
+
describe 'by index name' do
before do
allow(model).to receive(:index_exists_by_name?).with(:users, "index_x_by_y").and_return(true)
@@ -345,6 +365,16 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
model.remove_concurrent_index_by_name(:users, wrong_key: "index_x_by_y")
end.to raise_error 'remove_concurrent_index_by_name must get an index name as the second argument'
end
+
+ it 'unprepares the async index creation' do
+ expect(model).to receive(:remove_index)
+ .with(:users, { algorithm: :concurrently, name: "index_x_by_y" })
+
+ expect(model).to receive(:unprepare_async_index_by_name)
+ .with(:users, "index_x_by_y", { algorithm: :concurrently })
+
+ model.remove_concurrent_index_by_name(:users, "index_x_by_y")
+ end
end
end
end
@@ -384,9 +414,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
expect(model).to receive(:execute).with(/REFERENCES users \(id\)/)
@@ -398,9 +428,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
expect(model).to receive(:execute).with(/REFERENCES users \(id_convert_to_bigint\)/)
@@ -416,9 +446,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
expect(model).to receive(:execute).with(/ON DELETE SET NULL/)
@@ -433,9 +463,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
expect(model).to receive(:execute).with(/ON DELETE CASCADE/)
@@ -450,9 +480,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
expect(model).not_to receive(:execute).with(/ON DELETE/)
@@ -468,10 +498,10 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/NOT VALID/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
model.add_concurrent_foreign_key(:projects, :users, column: :user_id)
end
@@ -497,10 +527,10 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/NOT VALID/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT.+foo/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
model.add_concurrent_foreign_key(:projects, :users, column: :user_id, name: :foo)
end
@@ -527,10 +557,10 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/NOT VALID/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT.+bar/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
model.add_concurrent_foreign_key(:projects, :users, column: :user_id, name: :bar)
end
@@ -556,6 +586,22 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
it_behaves_like 'performs validation', {}
end
end
+
+ context 'when the reverse_lock_order flag is set' do
+ it 'explicitly locks the tables in target-source order', :aggregate_failures do
+ expect(model).to receive(:with_lock_retries).and_call_original
+ expect(model).to receive(:disable_statement_timeout).and_call_original
+ expect(model).to receive(:statement_timeout_disabled?).and_return(false)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
+ expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
+
+ expect(model).to receive(:execute).with('LOCK TABLE users, projects IN SHARE ROW EXCLUSIVE MODE')
+ expect(model).to receive(:execute).with(/REFERENCES users \(id\)/)
+
+ model.add_concurrent_foreign_key(:projects, :users, column: :user_id, reverse_lock_order: true)
+ end
+ end
end
end
@@ -568,9 +614,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).not_to receive(:concurrent_foreign_key_name)
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/ALTER TABLE projects VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
end
model.validate_foreign_key(:projects, :user_id, name: :foo)
@@ -585,9 +631,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:concurrent_foreign_key_name)
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(/ALTER TABLE projects VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
end
model.validate_foreign_key(:projects, :user_id)
@@ -702,7 +748,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
end
after do
- model.execute('RESET ALL')
+ model.execute('RESET statement_timeout')
end
it 'defines statement to 0 only for current transaction' do
@@ -719,7 +765,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
context 'when passing a blocks' do
it 'disables statement timeouts on session level and executes the block' do
expect(model).to receive(:execute).with('SET statement_timeout TO 0')
- expect(model).to receive(:execute).with('RESET ALL').at_least(:once)
+ expect(model).to receive(:execute).with('RESET statement_timeout').at_least(:once)
expect { |block| model.disable_statement_timeout(&block) }.to yield_control
end
@@ -731,7 +777,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
end
after do
- model.execute('RESET ALL')
+ model.execute('RESET statement_timeout')
end
it 'defines statement to 0 for any code run inside the block' do
@@ -758,12 +804,12 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
after do
# Use ActiveRecord::Base.connection instead of model.execute
# so that this call is not counted below
- ActiveRecord::Base.connection.execute('RESET ALL')
+ ActiveRecord::Base.connection.execute('RESET statement_timeout')
end
it 'yields control without disabling the timeout or resetting' do
expect(model).not_to receive(:execute).with('SET statement_timeout TO 0')
- expect(model).not_to receive(:execute).with('RESET ALL')
+ expect(model).not_to receive(:execute).with('RESET statement_timeout')
expect { |block| model.disable_statement_timeout(&block) }.to yield_control
end
@@ -2486,7 +2532,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
@@ -2496,7 +2542,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
.and_return(true).exactly(1)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
model.add_check_constraint(
:test_table,
@@ -2530,7 +2576,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:with_lock_retries).and_call_original
expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
@@ -2539,7 +2585,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
.and_return(true).exactly(1)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
model.add_check_constraint(
:test_table,
@@ -2572,9 +2618,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
expect(model).to receive(:check_constraint_exists?).and_return(true)
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:statement_timeout_disabled?).and_return(false)
- expect(model).to receive(:execute).with(/statement_timeout/)
+ expect(model).to receive(:execute).with(/SET statement_timeout TO/)
expect(model).to receive(:execute).ordered.with(validate_sql)
- expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+ expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
model.validate_check_constraint(:test_table, 'check_name')
end
diff --git a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
index e096e7f6e91..1a7116e75e5 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -581,4 +581,101 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
model.delete_queued_jobs('BackgroundMigrationClassName')
end
end
+
+ describe '#finalized_background_migration' do
+ include_context 'background migration job class'
+
+ let!(:tracked_pending_job) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1]) }
+ let!(:tracked_successful_job) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [2]) }
+
+ before do
+ Sidekiq::Testing.disable! do
+ BackgroundMigrationWorker.perform_async(job_class_name, [1, 2])
+ BackgroundMigrationWorker.perform_async(job_class_name, [3, 4])
+ BackgroundMigrationWorker.perform_in(10, job_class_name, [5, 6])
+ BackgroundMigrationWorker.perform_in(20, job_class_name, [7, 8])
+ end
+ end
+
+ it_behaves_like 'finalized tracked background migration' do
+ before do
+ model.finalize_background_migration(job_class_name)
+ end
+ end
+
+ context 'when removing all tracked job records' do
+ # Force pending jobs to remain pending.
+ let!(:job_perform_method) { ->(*arguments) { } }
+
+ before do
+ model.finalize_background_migration(job_class_name, delete_tracking_jobs: %w[pending succeeded])
+ end
+
+ it_behaves_like 'finalized tracked background migration'
+ it_behaves_like 'removed tracked jobs', 'pending'
+ it_behaves_like 'removed tracked jobs', 'succeeded'
+ end
+
+ context 'when retaining all tracked job records' do
+ before do
+ model.finalize_background_migration(job_class_name, delete_tracking_jobs: false)
+ end
+
+ it_behaves_like 'finalized background migration'
+ include_examples 'retained tracked jobs', 'succeeded'
+ end
+
+ context 'during retry race condition' do
+ let(:queue_items_added) { [] }
+ let!(:job_perform_method) do
+ ->(*arguments) do
+ Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
+ RSpec.current_example.example_group_instance.job_class_name,
+ arguments
+ )
+
+ # Mock another process pushing queue jobs.
+ queue_items_added = RSpec.current_example.example_group_instance.queue_items_added
+ if queue_items_added.count < 10
+ Sidekiq::Testing.disable! do
+ job_class_name = RSpec.current_example.example_group_instance.job_class_name
+ queue_items_added << BackgroundMigrationWorker.perform_async(job_class_name, [Time.current])
+ queue_items_added << BackgroundMigrationWorker.perform_in(10, job_class_name, [Time.current])
+ end
+ end
+ end
+ end
+
+ it_behaves_like 'finalized tracked background migration' do
+ before do
+ model.finalize_background_migration(job_class_name, delete_tracking_jobs: ['succeeded'])
+ end
+ end
+ end
+ end
+
+ describe '#delete_job_tracking' do
+ let!(:job_class_name) { 'TestJob' }
+
+ let!(:tracked_pending_job) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1]) }
+ let!(:tracked_successful_job) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [2]) }
+
+ context 'with default status' do
+ before do
+ model.delete_job_tracking(job_class_name)
+ end
+
+ include_examples 'retained tracked jobs', 'pending'
+ include_examples 'removed tracked jobs', 'succeeded'
+ end
+
+ context 'with explicit status' do
+ before do
+ model.delete_job_tracking(job_class_name, status: %w[pending succeeded])
+ end
+
+ include_examples 'removed tracked jobs', 'pending'
+ include_examples 'removed tracked jobs', 'succeeded'
+ end
+ end
end
diff --git a/spec/lib/gitlab/database/migrations/instrumentation_spec.rb b/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
index 6d047eed3bb..5945e5a2039 100644
--- a/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
+++ b/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
@@ -5,24 +5,35 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
describe '#observe' do
subject { described_class.new }
- let(:migration) { 1234 }
+ let(:migration_name) { 'test' }
+ let(:migration_version) { '12345' }
it 'executes the given block' do
- expect { |b| subject.observe(migration, &b) }.to yield_control
+ expect { |b| subject.observe(version: migration_version, name: migration_name, &b) }.to yield_control
end
context 'behavior with observers' do
- subject { described_class.new(observers).observe(migration) {} }
+ subject { described_class.new([Gitlab::Database::Migrations::Observers::MigrationObserver]).observe(version: migration_version, name: migration_name) {} }
- let(:observers) { [observer] }
let(:observer) { instance_double('Gitlab::Database::Migrations::Observers::MigrationObserver', before: nil, after: nil, record: nil) }
+ before do
+ allow(Gitlab::Database::Migrations::Observers::MigrationObserver).to receive(:new).and_return(observer)
+ end
+
+ it 'instantiates observer with observation' do
+ expect(Gitlab::Database::Migrations::Observers::MigrationObserver)
+ .to receive(:new)
+ .with(instance_of(Gitlab::Database::Migrations::Observation)) { |observation| expect(observation.version).to eq(migration_version) }
+ .and_return(observer)
+
+ subject
+ end
+
it 'calls #before, #after, #record on given observers' do
expect(observer).to receive(:before).ordered
expect(observer).to receive(:after).ordered
- expect(observer).to receive(:record).ordered do |observation|
- expect(observation.migration).to eq(migration)
- end
+ expect(observer).to receive(:record).ordered
subject
end
@@ -47,7 +58,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
end
context 'on successful execution' do
- subject { described_class.new.observe(migration) {} }
+ subject { described_class.new.observe(version: migration_version, name: migration_name) {} }
it 'records walltime' do
expect(subject.walltime).not_to be_nil
@@ -58,12 +69,16 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
end
it 'records the migration version' do
- expect(subject.migration).to eq(migration)
+ expect(subject.version).to eq(migration_version)
+ end
+
+ it 'records the migration name' do
+ expect(subject.name).to eq(migration_name)
end
end
context 'upon failure' do
- subject { described_class.new.observe(migration) { raise 'something went wrong' } }
+ subject { described_class.new.observe(version: migration_version, name: migration_name) { raise 'something went wrong' } }
it 'raises the exception' do
expect { subject }.to raise_error(/something went wrong/)
@@ -73,7 +88,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
subject { instance.observations.first }
before do
- instance.observe(migration) { raise 'something went wrong' }
+ instance.observe(version: migration_version, name: migration_name) { raise 'something went wrong' }
rescue StandardError
# ignore
end
@@ -89,7 +104,11 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
end
it 'records the migration version' do
- expect(subject.migration).to eq(migration)
+ expect(subject.version).to eq(migration_version)
+ end
+
+ it 'records the migration name' do
+ expect(subject.name).to eq(migration_name)
end
end
end
@@ -101,8 +120,8 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
let(:migration2) { double('migration2', call: nil) }
it 'records observations for all migrations' do
- subject.observe('migration1') {}
- subject.observe('migration2') { raise 'something went wrong' } rescue nil
+ subject.observe(version: migration_version, name: migration_name) {}
+ subject.observe(version: migration_version, name: migration_name) { raise 'something went wrong' } rescue nil
expect(subject.observations.size).to eq(2)
end
diff --git a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
index 8aac3ed67c6..36885a1594f 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
@@ -2,16 +2,17 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
- subject { described_class.new }
+ subject { described_class.new(observation) }
- let(:observation) { Gitlab::Database::Migrations::Observation.new(migration) }
+ let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
let(:connection) { ActiveRecord::Base.connection }
let(:query) { "select date_trunc('day', $1::timestamptz) + $2 * (interval '1 hour')" }
let(:query_binds) { [Time.current, 3] }
let(:directory_path) { Dir.mktmpdir }
- let(:log_file) { "#{directory_path}/#{migration}-query-details.json" }
+ let(:log_file) { "#{directory_path}/#{migration_version}_#{migration_name}-query-details.json" }
let(:query_details) { Gitlab::Json.parse(File.read(log_file)) }
- let(:migration) { 20210422152437 }
+ let(:migration_version) { 20210422152437 }
+ let(:migration_name) { 'test' }
before do
stub_const('Gitlab::Database::Migrations::Instrumentation::RESULT_DIR', directory_path)
@@ -49,7 +50,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
subject.before
run_query
subject.after
- subject.record(observation)
+ subject.record
end
def run_query
diff --git a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
index 195e7114582..2a49d8e8b73 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
@@ -2,14 +2,14 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
- subject { described_class.new }
+ subject { described_class.new(observation) }
- let(:observation) { Gitlab::Database::Migrations::Observation.new(migration) }
+ let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
let(:connection) { ActiveRecord::Base.connection }
let(:query) { 'select 1' }
let(:directory_path) { Dir.mktmpdir }
- let(:log_file) { "#{directory_path}/current.log" }
- let(:migration) { 20210422152437 }
+ let(:migration_version) { 20210422152437 }
+ let(:migration_name) { 'test' }
before do
stub_const('Gitlab::Database::Migrations::Instrumentation::RESULT_DIR', directory_path)
@@ -22,7 +22,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
it 'writes a file with the query log' do
observe
- expect(File.read("#{directory_path}/#{migration}.log")).to include(query)
+ expect(File.read("#{directory_path}/#{migration_version}_#{migration_name}.log")).to include(query)
end
it 'does not change the default logger' do
@@ -33,6 +33,6 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
subject.before
connection.execute(query)
subject.after
- subject.record(observation)
+ subject.record
end
end
diff --git a/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
index a3b03050b33..32a25fdaa28 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
@@ -2,8 +2,9 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::Migrations::Observers::QueryStatistics do
- subject { described_class.new }
+ subject { described_class.new(observation) }
+ let(:observation) { Gitlab::Database::Migrations::Observation.new }
let(:connection) { ActiveRecord::Base.connection }
def mock_pgss(enabled: true)
@@ -37,7 +38,6 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryStatistics do
end
describe '#record' do
- let(:observation) { Gitlab::Database::Migrations::Observation.new }
let(:result) { double }
let(:pgss_query) do
<<~SQL
@@ -52,7 +52,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryStatistics do
mock_pgss(enabled: true)
expect(connection).to receive(:execute).with(pgss_query).once.and_return(result)
- expect { subject.record(observation) }.to change { observation.query_statistics }.from(nil).to(result)
+ expect { subject.record }.to change { observation.query_statistics }.from(nil).to(result)
end
end
@@ -61,7 +61,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryStatistics do
mock_pgss(enabled: false)
expect(connection).not_to receive(:execute)
- expect { subject.record(observation) }.not_to change { observation.query_statistics }
+ expect { subject.record }.not_to change { observation.query_statistics }
end
end
end
diff --git a/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb b/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
index 73466471944..61e28003e66 100644
--- a/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
@@ -2,7 +2,7 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::Migrations::Observers::TotalDatabaseSizeChange do
- subject { described_class.new }
+ subject { described_class.new(observation) }
let(:observation) { Gitlab::Database::Migrations::Observation.new }
let(:connection) { ActiveRecord::Base.connection }
@@ -14,7 +14,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::TotalDatabaseSizeChange
subject.before
subject.after
- subject.record(observation)
+ subject.record
expect(observation.total_database_size_change).to eq(256 - 1024)
end
@@ -27,13 +27,13 @@ RSpec.describe Gitlab::Database::Migrations::Observers::TotalDatabaseSizeChange
it 'does not record anything if before size is unknown' do
subject.after
- expect { subject.record(observation) }.not_to change { observation.total_database_size_change }
+ expect { subject.record }.not_to change { observation.total_database_size_change }
end
it 'does not record anything if after size is unknown' do
subject.before
- expect { subject.record(observation) }.not_to change { observation.total_database_size_change }
+ expect { subject.record }.not_to change { observation.total_database_size_change }
end
end
end
diff --git a/spec/lib/gitlab/database/multi_threaded_migration_spec.rb b/spec/lib/gitlab/database/multi_threaded_migration_spec.rb
deleted file mode 100644
index 78dd9e88064..00000000000
--- a/spec/lib/gitlab/database/multi_threaded_migration_spec.rb
+++ /dev/null
@@ -1,43 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Gitlab::Database::MultiThreadedMigration do
- let(:migration) do
- Class.new { include Gitlab::Database::MultiThreadedMigration }.new
- end
-
- describe '#connection' do
- after do
- Thread.current[described_class::MULTI_THREAD_AR_CONNECTION] = nil
- end
-
- it 'returns the thread-local connection if present' do
- Thread.current[described_class::MULTI_THREAD_AR_CONNECTION] = 10
-
- expect(migration.connection).to eq(10)
- end
-
- it 'returns the global connection if no thread-local connection was set' do
- expect(migration.connection).to eq(ActiveRecord::Base.connection)
- end
- end
-
- describe '#with_multiple_threads' do
- it 'starts multiple threads and yields the supplied block in every thread' do
- output = Queue.new
-
- migration.with_multiple_threads(2) do
- output << migration.connection.execute('SELECT 1')
- end
-
- expect(output.size).to eq(2)
- end
-
- it 'joins the threads when the join parameter is set' do
- expect_any_instance_of(Thread).to receive(:join).and_call_original
-
- migration.with_multiple_threads(1) { }
- end
- end
-end
diff --git a/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb b/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
new file mode 100644
index 00000000000..8523b7104f0
--- /dev/null
+++ b/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
@@ -0,0 +1,181 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
+ include Database::TableSchemaHelpers
+
+ let(:connection) { ActiveRecord::Base.connection }
+
+ def expect_partition_present(name)
+ aggregate_failures do
+ expect(table_oid(name)).not_to be_nil
+ expect(Postgresql::DetachedPartition.find_by(table_name: name)).not_to be_nil
+ end
+ end
+
+ def expect_partition_removed(name)
+ aggregate_failures do
+ expect(table_oid(name)).to be_nil
+ expect(Postgresql::DetachedPartition.find_by(table_name: name)).to be_nil
+ end
+ end
+
+ before do
+ connection.execute(<<~SQL)
+ CREATE TABLE parent_table (
+ id bigserial not null,
+ created_at timestamptz not null,
+ primary key (id, created_at)
+ ) PARTITION BY RANGE(created_at)
+ SQL
+ end
+
+ def create_partition(name:, table: 'parent_table', from:, to:, attached:, drop_after:)
+ from = from.beginning_of_month
+ to = to.beginning_of_month
+ full_name = "#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.#{name}"
+ connection.execute(<<~SQL)
+ CREATE TABLE #{full_name}
+ PARTITION OF #{table}
+ FOR VALUES FROM ('#{from.strftime('%Y-%m-%d')}') TO ('#{to.strftime('%Y-%m-%d')}')
+ SQL
+
+ unless attached
+ connection.execute(<<~SQL)
+ ALTER TABLE #{table} DETACH PARTITION #{full_name}
+ SQL
+ end
+
+ Postgresql::DetachedPartition.create!(table_name: name,
+ drop_after: drop_after)
+ end
+
+ describe '#perform' do
+ context 'when the partition should not be dropped yet' do
+ it 'does not drop the partition' do
+ create_partition(name: 'test_partition',
+ from: 2.months.ago, to: 1.month.ago,
+ attached: false,
+ drop_after: 1.day.from_now)
+
+ subject.perform
+
+ expect_partition_present('test_partition')
+ end
+ end
+
+ context 'with a partition to drop' do
+ before do
+ create_partition(name: 'test_partition',
+ from: 2.months.ago,
+ to: 1.month.ago.beginning_of_month,
+ attached: false,
+ drop_after: 1.second.ago)
+ end
+
+ it 'drops the partition' do
+ subject.perform
+
+ expect(table_oid('test_partition')).to be_nil
+ end
+
+ context 'when the drop_detached_partitions feature flag is disabled' do
+ before do
+ stub_feature_flags(drop_detached_partitions: false)
+ end
+ it 'does not drop the partition' do
+ subject.perform
+
+ expect(table_oid('test_partition')).not_to be_nil
+ end
+ end
+
+ context 'when another process drops the table while the first waits for a lock' do
+ it 'skips the table' do
+ # Rspec's receive_method_chain does not support .and_wrap_original, so we need to nest here.
+ expect(Postgresql::DetachedPartition).to receive(:lock).and_wrap_original do |lock_meth|
+ locked = lock_meth.call
+ expect(locked).to receive(:find_by).and_wrap_original do |find_meth, *find_args|
+ # Another process drops the table then deletes this entry
+ Postgresql::DetachedPartition.where(*find_args).delete_all
+ find_meth.call(*find_args)
+ end
+
+ locked
+ end
+
+ expect(subject).not_to receive(:drop_one)
+
+ subject.perform
+ end
+ end
+ end
+
+ context 'when the partition to drop is still attached to its table' do
+ before do
+ create_partition(name: 'test_partition',
+ from: 2.months.ago,
+ to: 1.month.ago.beginning_of_month,
+ attached: true,
+ drop_after: 1.second.ago)
+ end
+
+ it 'does not drop the partition, but does remove the DetachedPartition entry' do
+ subject.perform
+ aggregate_failures do
+ expect(table_oid('test_partition')).not_to be_nil
+ expect(Postgresql::DetachedPartition.find_by(table_name: 'test_partition')).to be_nil
+ end
+ end
+
+ it 'removes the detached_partition entry' do
+ detached_partition = Postgresql::DetachedPartition.find_by!(table_name: 'test_partition')
+
+ subject.perform
+
+ expect(Postgresql::DetachedPartition.exists?(id: detached_partition.id)).to be_falsey
+ end
+ end
+
+ context 'with multiple partitions to drop' do
+ before do
+ create_partition(name: 'partition_1',
+ from: 3.months.ago,
+ to: 2.months.ago,
+ attached: false,
+ drop_after: 1.second.ago)
+
+ create_partition(name: 'partition_2',
+ from: 2.months.ago,
+ to: 1.month.ago,
+ attached: false,
+ drop_after: 1.second.ago)
+ end
+
+ it 'drops both partitions' do
+ subject.perform
+
+ expect_partition_removed('partition_1')
+ expect_partition_removed('partition_2')
+ end
+
+ context 'when the first drop returns an error' do
+ it 'still drops the second partition' do
+ expect(subject).to receive(:drop_one).ordered.and_raise('injected error')
+ expect(subject).to receive(:drop_one).ordered.and_call_original
+
+ subject.perform
+
+ # We don't know which partition we tried to drop first, so the tests here have to work with either one
+ expect(Postgresql::DetachedPartition.count).to eq(1)
+ errored_partition_name = Postgresql::DetachedPartition.first!.table_name
+
+ dropped_partition_name = (%w[partition_1 partition_2] - [errored_partition_name]).first
+ expect_partition_present(errored_partition_name)
+ expect_partition_removed(dropped_partition_name)
+ end
+ end
+ end
+ end
+end
diff --git a/spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb b/spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb
index f9dca371398..c4fbf53d1c2 100644
--- a/spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb
@@ -237,16 +237,6 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
expect(subject).to contain_exactly(min_value_to_may)
end
-
- context 'when the feature flag is toggled off' do
- before do
- stub_feature_flags(partition_pruning_dry_run: false)
- end
-
- it 'is empty' do
- expect(subject).to eq([])
- end
- end
end
context 'with a time retention policy of 2 months' do
@@ -258,16 +248,6 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-05-01', '2020-06-01', partition_name: 'partitioned_test_202005')
)
end
-
- context 'when the feature flag is toggled off' do
- before do
- stub_feature_flags(partition_pruning_dry_run: false)
- end
-
- it 'is empty' do
- expect(subject).to eq([])
- end
- end
end
end
end
diff --git a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
index 903a41d6dd2..3d60457c3a9 100644
--- a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
@@ -4,9 +4,14 @@ require 'spec_helper'
RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
include Database::PartitioningHelpers
- include Database::TableSchemaHelpers
include ExclusiveLeaseHelpers
+ def has_partition(model, month)
+ Gitlab::Database::PostgresPartition.for_parent_table(model.table_name).any? do |partition|
+ Gitlab::Database::Partitioning::TimePartition.from_sql(model.table_name, partition.name, partition.condition).from == month
+ end
+ end
+
describe '.register' do
let(:model) { double(partitioning_strategy: nil) }
@@ -111,14 +116,14 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
let(:extra_partitions) do
[
- instance_double(Gitlab::Database::Partitioning::TimePartition, table: table, partition_name: 'foo1'),
- instance_double(Gitlab::Database::Partitioning::TimePartition, table: table, partition_name: 'foo2')
+ instance_double(Gitlab::Database::Partitioning::TimePartition, table: table, partition_name: 'foo1', to_detach_sql: 'SELECT 1'),
+ instance_double(Gitlab::Database::Partitioning::TimePartition, table: table, partition_name: 'foo2', to_detach_sql: 'SELECT 2')
]
end
- context 'with the partition_pruning_dry_run feature flag enabled' do
+ context 'with the partition_pruning feature flag enabled' do
before do
- stub_feature_flags(partition_pruning_dry_run: true)
+ stub_feature_flags(partition_pruning: true)
end
it 'detaches each extra partition' do
@@ -146,9 +151,9 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
end
end
- context 'with the partition_pruning_dry_run feature flag disabled' do
+ context 'with the partition_pruning feature flag disabled' do
before do
- stub_feature_flags(partition_pruning_dry_run: false)
+ stub_feature_flags(partition_pruning: false)
end
it 'returns immediately' do
@@ -158,4 +163,128 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
end
end
end
+
+ describe '#detach_partitions' do
+ around do |ex|
+ travel_to(Date.parse('2021-06-23')) do
+ ex.run
+ end
+ end
+
+ subject { described_class.new([my_model]).sync_partitions }
+
+ let(:connection) { ActiveRecord::Base.connection }
+ let(:my_model) do
+ Class.new(ApplicationRecord) do
+ include PartitionedTable
+
+ self.table_name = 'my_model_example_table'
+
+ partitioned_by :created_at, strategy: :monthly, retain_for: 1.month
+ end
+ end
+
+ before do
+ connection.execute(<<~SQL)
+ CREATE TABLE my_model_example_table
+ (id serial not null, created_at timestamptz not null, primary key (id, created_at))
+ PARTITION BY RANGE (created_at);
+
+ CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.my_model_example_table_202104
+ PARTITION OF my_model_example_table
+ FOR VALUES FROM ('2021-04-01') TO ('2021-05-01');
+
+ CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.my_model_example_table_202105
+ PARTITION OF my_model_example_table
+ FOR VALUES FROM ('2021-05-01') TO ('2021-06-01');
+ SQL
+
+ # Also create all future partitions so that the sync is only trying to detach old partitions
+ my_model.partitioning_strategy.missing_partitions.each do |p|
+ connection.execute p.to_sql
+ end
+ end
+
+ def num_tables
+ connection.select_value(<<~SQL)
+ SELECT COUNT(*)
+ FROM pg_class
+ where relkind IN ('r', 'p')
+ SQL
+ end
+
+ it 'detaches exactly one partition' do
+ expect { subject }.to change { find_partitions(my_model.table_name, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA).size }.from(9).to(8)
+ end
+
+ it 'detaches the old partition' do
+ expect { subject }.to change { has_partition(my_model, 2.months.ago.beginning_of_month) }.from(true).to(false)
+ end
+
+ it 'deletes zero tables' do
+ expect { subject }.not_to change { num_tables }
+ end
+
+ it 'creates the appropriate PendingPartitionDrop entry' do
+ subject
+
+ pending_drop = Postgresql::DetachedPartition.find_by!(table_name: 'my_model_example_table_202104')
+ expect(pending_drop.drop_after).to eq(Time.current + described_class::RETAIN_DETACHED_PARTITIONS_FOR)
+ end
+
+ # Postgres 11 does not support foreign keys to partitioned tables
+ if Gitlab::Database.main.version.to_f >= 12
+ context 'when the model is the target of a foreign key' do
+ before do
+ connection.execute(<<~SQL)
+ create unique index idx_for_fk ON my_model_example_table(created_at);
+
+ create table referencing_table (
+ id bigserial primary key not null,
+ referencing_created_at timestamptz references my_model_example_table(created_at)
+ );
+ SQL
+ end
+
+ it 'does not detach partitions with a referenced foreign key' do
+ expect { subject }.not_to change { find_partitions(my_model.table_name).size }
+ end
+ end
+ end
+ end
+
+ context 'creating and then detaching partitions for a table' do
+ let(:connection) { ActiveRecord::Base.connection }
+ let(:my_model) do
+ Class.new(ApplicationRecord) do
+ include PartitionedTable
+
+ self.table_name = 'my_model_example_table'
+
+ partitioned_by :created_at, strategy: :monthly, retain_for: 1.month
+ end
+ end
+
+ before do
+ connection.execute(<<~SQL)
+ CREATE TABLE my_model_example_table
+ (id serial not null, created_at timestamptz not null, primary key (id, created_at))
+ PARTITION BY RANGE (created_at);
+ SQL
+ end
+
+ def num_partitions(model)
+ find_partitions(model.table_name, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA).size
+ end
+
+ it 'creates partitions for the future then drops the oldest one after a month' do
+ # 1 month for the current month, 1 month for the old month that we're retaining data for, headroom
+ expected_num_partitions = (Gitlab::Database::Partitioning::MonthlyStrategy::HEADROOM + 2.months) / 1.month
+ expect { described_class.new([my_model]).sync_partitions }.to change { num_partitions(my_model) }.from(0).to(expected_num_partitions)
+
+ travel 1.month
+
+ expect { described_class.new([my_model]).sync_partitions }.to change { has_partition(my_model, 2.months.ago.beginning_of_month) }.from(true).to(false).and(change { num_partitions(my_model) }.by(0))
+ end
+ end
end
diff --git a/spec/lib/gitlab/database/partitioning/partition_monitoring_spec.rb b/spec/lib/gitlab/database/partitioning/partition_monitoring_spec.rb
index 67596211f71..7024cbd55ff 100644
--- a/spec/lib/gitlab/database/partitioning/partition_monitoring_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/partition_monitoring_spec.rb
@@ -8,7 +8,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionMonitoring do
let(:models) { [model] }
let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table) }
- let(:partitioning_strategy) { double(missing_partitions: missing_partitions, current_partitions: current_partitions) }
+ let(:partitioning_strategy) { double(missing_partitions: missing_partitions, current_partitions: current_partitions, extra_partitions: extra_partitions) }
let(:table) { "some_table" }
let(:missing_partitions) do
@@ -19,6 +19,10 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionMonitoring do
[double, double]
end
+ let(:extra_partitions) do
+ [double, double, double]
+ end
+
it 'reports number of present partitions' do
subject
@@ -30,5 +34,11 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionMonitoring do
expect(Gitlab::Metrics.registry.get(:db_partitions_missing).get({ table: table })).to eq(missing_partitions.size)
end
+
+ it 'reports number of extra partitions' do
+ subject
+
+ expect(Gitlab::Metrics.registry.get(:db_partitions_extra).get({ table: table })).to eq(extra_partitions.size)
+ end
end
end
diff --git a/spec/lib/gitlab/database/postgres_foreign_key_spec.rb b/spec/lib/gitlab/database/postgres_foreign_key_spec.rb
new file mode 100644
index 00000000000..ec39e5bfee7
--- /dev/null
+++ b/spec/lib/gitlab/database/postgres_foreign_key_spec.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model do
+ # PostgresForeignKey does not `behaves_like 'a postgres model'` because it does not correspond 1-1 with a single entry
+ # in pg_class
+
+ before do
+ ActiveRecord::Base.connection.execute(<<~SQL)
+ CREATE TABLE public.referenced_table (
+ id bigserial primary key not null
+ );
+
+ CREATE TABLE public.other_referenced_table (
+ id bigserial primary key not null
+ );
+
+ CREATE TABLE public.constrained_table (
+ id bigserial primary key not null,
+ referenced_table_id bigint not null,
+ other_referenced_table_id bigint not null,
+ CONSTRAINT fk_constrained_to_referenced FOREIGN KEY(referenced_table_id) REFERENCES referenced_table(id),
+ CONSTRAINT fk_constrained_to_other_referenced FOREIGN KEY(other_referenced_table_id)
+ REFERENCES other_referenced_table(id)
+ );
+ SQL
+ end
+
+ describe '#by_referenced_table_identifier' do
+ it 'throws an error when the identifier name is not fully qualified' do
+ expect { described_class.by_referenced_table_identifier('referenced_table') }.to raise_error(ArgumentError, /not fully qualified/)
+ end
+
+ it 'finds the foreign keys for the referenced table' do
+ expected = described_class.find_by!(name: 'fk_constrained_to_referenced')
+
+ expect(described_class.by_referenced_table_identifier('public.referenced_table')).to contain_exactly(expected)
+ end
+ end
+end
diff --git a/spec/lib/gitlab/database/postgres_index_spec.rb b/spec/lib/gitlab/database/postgres_index_spec.rb
index e1832219ebf..9088719d5a4 100644
--- a/spec/lib/gitlab/database/postgres_index_spec.rb
+++ b/spec/lib/gitlab/database/postgres_index_spec.rb
@@ -40,6 +40,37 @@ RSpec.describe Gitlab::Database::PostgresIndex do
expect(types & %w(btree gist)).to eq(types)
end
+
+ context 'with leftover indexes' do
+ before do
+ ActiveRecord::Base.connection.execute(<<~SQL)
+ CREATE INDEX foobar_ccnew ON users (id);
+ CREATE INDEX foobar_ccnew1 ON users (id);
+ SQL
+ end
+
+ subject { described_class.reindexing_support.map(&:name) }
+
+ it 'excludes temporary indexes from reindexing' do
+ expect(subject).not_to include('foobar_ccnew')
+ expect(subject).not_to include('foobar_ccnew1')
+ end
+ end
+ end
+
+ describe '.reindexing_leftovers' do
+ subject { described_class.reindexing_leftovers }
+
+ before do
+ ActiveRecord::Base.connection.execute(<<~SQL)
+ CREATE INDEX foobar_ccnew ON users (id);
+ CREATE INDEX foobar_ccnew1 ON users (id);
+ SQL
+ end
+
+ it 'retrieves leftover indexes matching the /_ccnew[0-9]*$/ pattern' do
+ expect(subject.map(&:name)).to eq(%w(foobar_ccnew foobar_ccnew1))
+ end
end
describe '.not_match' do
diff --git a/spec/lib/gitlab/database/reindexing_spec.rb b/spec/lib/gitlab/database/reindexing_spec.rb
index 8aff99544ca..550f9db2b5b 100644
--- a/spec/lib/gitlab/database/reindexing_spec.rb
+++ b/spec/lib/gitlab/database/reindexing_spec.rb
@@ -26,14 +26,31 @@ RSpec.describe Gitlab::Database::Reindexing do
end
end
- describe '.candidate_indexes' do
- subject { described_class.candidate_indexes }
+ describe '.cleanup_leftovers!' do
+ subject { described_class.cleanup_leftovers! }
+
+ before do
+ ApplicationRecord.connection.execute(<<~SQL)
+ CREATE INDEX foobar_ccnew ON users (id);
+ CREATE INDEX foobar_ccnew1 ON users (id);
+ SQL
+ end
- it 'retrieves regular indexes that are no left-overs from previous runs' do
- result = double
- expect(Gitlab::Database::PostgresIndex).to receive_message_chain('not_match.reindexing_support').with('\_ccnew[0-9]*$').with(no_args).and_return(result)
+ it 'drops both leftover indexes' do
+ expect_query("SET lock_timeout TO '60000ms'")
+ expect_query("DROP INDEX CONCURRENTLY IF EXISTS \"public\".\"foobar_ccnew\"")
+ expect_query("RESET idle_in_transaction_session_timeout; RESET lock_timeout")
+ expect_query("SET lock_timeout TO '60000ms'")
+ expect_query("DROP INDEX CONCURRENTLY IF EXISTS \"public\".\"foobar_ccnew1\"")
+ expect_query("RESET idle_in_transaction_session_timeout; RESET lock_timeout")
- expect(subject).to eq(result)
+ subject
+ end
+
+ def expect_query(sql)
+ expect(ApplicationRecord.connection).to receive(:execute).ordered.with(sql).and_wrap_original do |method, sql|
+ method.call(sql.sub(/CONCURRENTLY/, ''))
+ end
end
end
end
diff --git a/spec/lib/gitlab/database/schema_migrations/context_spec.rb b/spec/lib/gitlab/database/schema_migrations/context_spec.rb
index f3bed9b40d6..1f1943d00a3 100644
--- a/spec/lib/gitlab/database/schema_migrations/context_spec.rb
+++ b/spec/lib/gitlab/database/schema_migrations/context_spec.rb
@@ -3,7 +3,8 @@
require 'spec_helper'
RSpec.describe Gitlab::Database::SchemaMigrations::Context do
- let(:connection) { ActiveRecord::Base.connection }
+ let(:connection_class) { ActiveRecord::Base }
+ let(:connection) { connection_class.connection }
let(:context) { described_class.new(connection) }
@@ -12,13 +13,65 @@ RSpec.describe Gitlab::Database::SchemaMigrations::Context do
expect(context.schema_directory).to eq(File.join(Rails.root, 'db/schema_migrations'))
end
- context 'multiple databases' do
- let(:connection) { Ci::BaseModel.connection }
+ context 'CI database' do
+ let(:connection_class) { Ci::CiDatabaseRecord }
it 'returns a directory path that is database specific' do
skip_if_multiple_databases_not_setup
- expect(context.schema_directory).to eq(File.join(Rails.root, 'db/ci_schema_migrations'))
+ expect(context.schema_directory).to eq(File.join(Rails.root, 'db/schema_migrations'))
+ end
+ end
+
+ context 'multiple databases' do
+ let(:connection_class) do
+ Class.new(::ApplicationRecord) do
+ self.abstract_class = true
+
+ def self.name
+ 'Gitlab::Database::SchemaMigrations::Context::TestConnection'
+ end
+ end
+ end
+
+ let(:configuration_overrides) { {} }
+
+ before do
+ connection_class.establish_connection(
+ ActiveRecord::Base
+ .connection_pool
+ .db_config
+ .configuration_hash
+ .merge(configuration_overrides)
+ )
+ end
+
+ after do
+ connection_class.remove_connection
+ end
+
+ context 'when `schema_migrations_path` is configured as string' do
+ let(:configuration_overrides) do
+ { "schema_migrations_path" => "db/ci_schema_migrations" }
+ end
+
+ it 'returns a configured directory path that' do
+ skip_if_multiple_databases_not_setup
+
+ expect(context.schema_directory).to eq(File.join(Rails.root, 'db/ci_schema_migrations'))
+ end
+ end
+
+ context 'when `schema_migrations_path` is configured as symbol' do
+ let(:configuration_overrides) do
+ { schema_migrations_path: "db/ci_schema_migrations" }
+ end
+
+ it 'returns a configured directory path that' do
+ skip_if_multiple_databases_not_setup
+
+ expect(context.schema_directory).to eq(File.join(Rails.root, 'db/ci_schema_migrations'))
+ end
end
end
end
diff --git a/spec/lib/gitlab/database/transaction/context_spec.rb b/spec/lib/gitlab/database/transaction/context_spec.rb
new file mode 100644
index 00000000000..65d52b4d099
--- /dev/null
+++ b/spec/lib/gitlab/database/transaction/context_spec.rb
@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Transaction::Context do
+ subject { described_class.new }
+
+ let(:data) { subject.context }
+
+ before do
+ stub_const("#{described_class}::LOG_THROTTLE", 100)
+ end
+
+ describe '#set_start_time' do
+ before do
+ subject.set_start_time
+ end
+
+ it 'sets start_time' do
+ expect(data).to have_key(:start_time)
+ end
+ end
+
+ describe '#increment_savepoints' do
+ before do
+ 2.times { subject.increment_savepoints }
+ end
+
+ it { expect(data[:savepoints]).to eq(2) }
+ end
+
+ describe '#increment_rollbacks' do
+ before do
+ 3.times { subject.increment_rollbacks }
+ end
+
+ it { expect(data[:rollbacks]).to eq(3) }
+ end
+
+ describe '#increment_releases' do
+ before do
+ 4.times { subject.increment_releases }
+ end
+
+ it { expect(data[:releases]).to eq(4) }
+ end
+
+ describe '#set_depth' do
+ before do
+ subject.set_depth(2)
+ end
+
+ it { expect(data[:depth]).to eq(2) }
+ end
+
+ describe '#track_sql' do
+ before do
+ subject.track_sql('SELECT 1')
+ subject.track_sql('SELECT * FROM users')
+ end
+
+ it { expect(data[:queries]).to eq(['SELECT 1', 'SELECT * FROM users']) }
+ end
+
+ describe '#duration' do
+ before do
+ subject.set_start_time
+ end
+
+ it { expect(subject.duration).to be >= 0 }
+ end
+
+ context 'when depth is low' do
+ it 'does not log data upon COMMIT' do
+ expect(subject).not_to receive(:application_info)
+
+ subject.commit
+ end
+
+ it 'does not log data upon ROLLBACK' do
+ expect(subject).not_to receive(:application_info)
+
+ subject.rollback
+ end
+
+ it '#should_log? returns false' do
+ expect(subject.should_log?).to be false
+ end
+ end
+
+ shared_examples 'logs transaction data' do
+ it 'logs once upon COMMIT' do
+ expect(subject).to receive(:application_info).and_call_original
+
+ 2.times { subject.commit }
+ end
+
+ it 'logs once upon ROLLBACK' do
+ expect(subject).to receive(:application_info).once
+
+ 2.times { subject.rollback }
+ end
+
+ it 'logs again when log throttle duration passes' do
+ expect(subject).to receive(:application_info).twice.and_call_original
+
+ 2.times { subject.commit }
+
+ data[:last_log_timestamp] -= (described_class::LOG_THROTTLE_DURATION + 1)
+
+ subject.commit
+ end
+
+ it '#should_log? returns true' do
+ expect(subject.should_log?).to be true
+ end
+ end
+
+ context 'when depth exceeds threshold' do
+ before do
+ subject.set_depth(described_class::LOG_DEPTH_THRESHOLD + 1)
+ end
+
+ it_behaves_like 'logs transaction data'
+ end
+
+ context 'when savepoints count exceeds threshold' do
+ before do
+ data[:savepoints] = described_class::LOG_SAVEPOINTS_THRESHOLD + 1
+ end
+
+ it_behaves_like 'logs transaction data'
+ end
+
+ context 'when duration exceeds threshold' do
+ before do
+ subject.set_start_time
+
+ data[:start_time] -= (described_class::LOG_DURATION_S_THRESHOLD + 1)
+ end
+
+ it_behaves_like 'logs transaction data'
+ end
+end
diff --git a/spec/lib/gitlab/database/transaction/observer_spec.rb b/spec/lib/gitlab/database/transaction/observer_spec.rb
new file mode 100644
index 00000000000..7aa24217dc3
--- /dev/null
+++ b/spec/lib/gitlab/database/transaction/observer_spec.rb
@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Transaction::Observer do
+ # Use the delete DB strategy so that the test won't be wrapped in a transaction
+ describe '.instrument_transactions', :delete do
+ let(:transaction_context) { ActiveRecord::Base.connection.transaction_manager.transaction_context }
+ let(:context) { transaction_context.context }
+
+ around do |example|
+ # Emulate production environment when SQL comments come first to avoid truncation
+ Marginalia::Comment.prepend_comment = true
+ subscriber = described_class.register!
+
+ example.run
+
+ ActiveSupport::Notifications.unsubscribe(subscriber)
+ Marginalia::Comment.prepend_comment = false
+ end
+
+ it 'tracks transaction data', :aggregate_failures do
+ ActiveRecord::Base.transaction do
+ ActiveRecord::Base.transaction(requires_new: true) do
+ User.first
+
+ expect(transaction_context).to be_a(::Gitlab::Database::Transaction::Context)
+ expect(context.keys).to match_array(%i(start_time depth savepoints queries))
+ expect(context[:depth]).to eq(2)
+ expect(context[:savepoints]).to eq(1)
+ expect(context[:queries].length).to eq(1)
+ end
+ end
+
+ expect(context[:depth]).to eq(2)
+ expect(context[:savepoints]).to eq(1)
+ expect(context[:releases]).to eq(1)
+ end
+
+ describe '.extract_sql_command' do
+ using RSpec::Parameterized::TableSyntax
+
+ where(:sql, :expected) do
+ 'SELECT 1' | 'SELECT 1'
+ '/* test comment */ SELECT 1' | 'SELECT 1'
+ '/* test comment */ ROLLBACK TO SAVEPOINT point1' | 'ROLLBACK TO SAVEPOINT '
+ 'SELECT 1 /* trailing comment */' | 'SELECT 1 /* trailing comment */'
+ end
+
+ with_them do
+ it do
+ expect(described_class.extract_sql_command(sql)).to eq(expected)
+ end
+ end
+ end
+ end
+end
diff --git a/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb b/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
index ff8e76311ae..0282a7af0df 100644
--- a/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
+++ b/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
@@ -37,22 +37,20 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
context 'when lock retry is enabled' do
let(:lock_fiber) do
Fiber.new do
- configuration = ActiveRecordSecond.configurations.find_db_config(Rails.env).configuration_hash
+ # Initiating a separate DB connection for the lock
+ conn = ActiveRecord::Base.connection_pool.checkout
- # Initiating a second DB connection for the lock
- conn = ActiveRecordSecond.establish_connection(configuration).connection
conn.transaction do
conn.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
Fiber.yield
end
- ActiveRecordSecond.remove_connection # force disconnect
+ # Releasing the connection we requested
+ ActiveRecord::Base.connection_pool.checkin(conn)
end
end
before do
- stub_const('ActiveRecordSecond', Class.new(ActiveRecord::Base))
-
lock_fiber.resume # start the transaction and lock the table
end
diff --git a/spec/lib/gitlab/database/with_lock_retries_spec.rb b/spec/lib/gitlab/database/with_lock_retries_spec.rb
index 367f793b117..72074f06210 100644
--- a/spec/lib/gitlab/database/with_lock_retries_spec.rb
+++ b/spec/lib/gitlab/database/with_lock_retries_spec.rb
@@ -37,22 +37,19 @@ RSpec.describe Gitlab::Database::WithLockRetries do
context 'when lock retry is enabled' do
let(:lock_fiber) do
Fiber.new do
- configuration = ActiveRecordSecond.configurations.find_db_config(Rails.env).configuration_hash
-
- # Initiating a second DB connection for the lock
- conn = ActiveRecordSecond.establish_connection(configuration).connection
+ # Initiating a separate DB connection for the lock
+ conn = ActiveRecord::Base.connection_pool.checkout
conn.transaction do
conn.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
Fiber.yield
end
- ActiveRecordSecond.remove_connection # force disconnect
+ # Releasing the connection we requested
+ ActiveRecord::Base.connection_pool.checkin(conn)
end
end
before do
- stub_const('ActiveRecordSecond', Class.new(ActiveRecord::Base))
-
lock_fiber.resume # start the transaction and lock the table
end