diff options
Diffstat (limited to 'spec/lib/gitlab/database/partitioning/partition_manager_spec.rb')
-rw-r--r-- | spec/lib/gitlab/database/partitioning/partition_manager_spec.rb | 161 |
1 files changed, 161 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb new file mode 100644 index 00000000000..903a41d6dd2 --- /dev/null +++ b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb @@ -0,0 +1,161 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::Partitioning::PartitionManager do + include Database::PartitioningHelpers + include Database::TableSchemaHelpers + include ExclusiveLeaseHelpers + + describe '.register' do + let(:model) { double(partitioning_strategy: nil) } + + it 'remembers registered models' do + expect { described_class.register(model) }.to change { described_class.models }.to include(model) + end + end + + context 'creating partitions (mocked)' do + subject(:sync_partitions) { described_class.new(models).sync_partitions } + + let(:models) { [model] } + let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table) } + let(:partitioning_strategy) { double(missing_partitions: partitions, extra_partitions: []) } + let(:table) { "some_table" } + + before do + allow(ActiveRecord::Base.connection).to receive(:table_exists?).and_call_original + allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true) + allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original + + stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT) + end + + let(:partitions) do + [ + instance_double(Gitlab::Database::Partitioning::TimePartition, table: 'bar', partition_name: 'foo', to_sql: "SELECT 1"), + instance_double(Gitlab::Database::Partitioning::TimePartition, table: 'bar', partition_name: 'foo2', to_sql: "SELECT 2") + ] + end + + it 'creates the partition' do + expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.first.to_sql) + expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.second.to_sql) + + sync_partitions + end + + context 'error handling with 2 models' do + let(:models) do + [ + double(partitioning_strategy: strategy1, table_name: table), + double(partitioning_strategy: strategy2, table_name: table) + ] + end + + let(:strategy1) { double('strategy1', missing_partitions: nil, extra_partitions: []) } + let(:strategy2) { double('strategy2', missing_partitions: partitions, extra_partitions: []) } + + it 'still creates partitions for the second table' do + expect(strategy1).to receive(:missing_partitions).and_raise('this should never happen (tm)') + expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.first.to_sql) + expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.second.to_sql) + + sync_partitions + end + end + end + + context 'creating partitions' do + subject(:sync_partitions) { 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 + 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 + + it 'creates partitions' do + expect { sync_partitions }.to change { find_partitions(my_model.table_name, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA).size }.from(0) + end + end + + context 'detaching partitions (mocked)' do + subject(:sync_partitions) { manager.sync_partitions } + + let(:manager) { described_class.new(models) } + let(:models) { [model] } + let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table)} + let(:partitioning_strategy) { double(extra_partitions: extra_partitions, missing_partitions: []) } + let(:table) { "foo" } + + before do + allow(ActiveRecord::Base.connection).to receive(:table_exists?).and_call_original + allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true) + + stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT) + end + + 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') + ] + end + + context 'with the partition_pruning_dry_run feature flag enabled' do + before do + stub_feature_flags(partition_pruning_dry_run: true) + end + + it 'detaches each extra partition' do + extra_partitions.each { |p| expect(manager).to receive(:detach_one_partition).with(p) } + + sync_partitions + end + + context 'error handling' do + let(:models) do + [ + double(partitioning_strategy: error_strategy, table_name: table), + model + ] + end + + let(:error_strategy) { double(extra_partitions: nil, missing_partitions: []) } + + it 'still drops partitions for the other model' do + expect(error_strategy).to receive(:extra_partitions).and_raise('injected error!') + extra_partitions.each { |p| expect(manager).to receive(:detach_one_partition).with(p) } + + sync_partitions + end + end + end + + context 'with the partition_pruning_dry_run feature flag disabled' do + before do + stub_feature_flags(partition_pruning_dry_run: false) + end + + it 'returns immediately' do + expect(manager).not_to receive(:detach) + + sync_partitions + end + end + end +end |