1 files changed, 121 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb b/spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb
index 885eef5723e..f9dca371398 100644
--- a/spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb
@@ -71,6 +71,18 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
         model.create!(created_at: Date.parse('2020-06-15'))
       end
 
+      context 'when pruning partitions before June 2020' do
+        subject { described_class.new(model, partitioning_key, retain_for: 1.month).missing_partitions }
+
+        it 'does not include the missing partition from May 2020 because it would be dropped' do
+          expect(subject).not_to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-05-01', '2020-06-01'))
+        end
+
+        it 'detects the missing partition for 1 month ago (July 2020)' do
+          expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-07-01', '2020-08-01'))
+        end
+      end
+
       it 'detects the gap and the missing partition in May 2020' do
         expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-05-01', '2020-06-01'))
       end
@@ -108,6 +120,19 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
         SQL
       end
 
+      context 'when pruning partitions before June 2020' do
+        subject { described_class.new(model, partitioning_key, retain_for: 1.month).missing_partitions }
+
+        it 'detects exactly the set of partitions from June 2020 to March 2021' do
+          months = %w[2020-07-01 2020-08-01 2020-09-01 2020-10-01 2020-11-01 2020-12-01 2021-01-01 2021-02-01 2021-03-01]
+          expected = months[..-2].zip(months.drop(1)).map do |(from, to)|
+            Gitlab::Database::Partitioning::TimePartition.new(model.table_name, from, to)
+          end
+
+          expect(subject).to match_array(expected)
+        end
+      end
+
       it 'detects the missing catch-all partition at the beginning' do
         expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, nil, '2020-08-01'))
       end
@@ -150,4 +175,100 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
       end
     end
   end
+
+  describe '#extra_partitions' do
+    let(:model) do
+      Class.new(ActiveRecord::Base) do
+        self.table_name = 'partitioned_test'
+        self.primary_key = :id
+      end
+    end
+
+    let(:partitioning_key) { :created_at }
+    let(:table_name) { :partitioned_test }
+
+    around do |example|
+      travel_to(Date.parse('2020-08-22')) { example.run }
+    end
+
+    describe 'with existing partitions' do
+      before do
+        ActiveRecord::Base.connection.execute(<<~SQL)
+          CREATE TABLE #{table_name}
+            (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}.partitioned_test_000000
+          PARTITION OF #{table_name}
+          FOR VALUES FROM (MINVALUE) TO ('2020-05-01');
+
+          CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.partitioned_test_202005
+          PARTITION OF #{table_name}
+          FOR VALUES FROM ('2020-05-01') TO ('2020-06-01');
+
+          CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.partitioned_test_202006
+          PARTITION OF #{table_name}
+          FOR VALUES FROM ('2020-06-01') TO ('2020-07-01')
+        SQL
+      end
+
+      context 'without a time retention policy' do
+        subject { described_class.new(model, partitioning_key).extra_partitions }
+
+        it 'has no extra partitions to prune' do
+          expect(subject).to eq([])
+        end
+      end
+
+      context 'with a time retention policy that excludes no partitions' do
+        subject { described_class.new(model, partitioning_key, retain_for: 4.months).extra_partitions }
+
+        it 'has no extra partitions to prune' do
+          expect(subject).to eq([])
+        end
+      end
+
+      context 'with a time retention policy of 3 months' do
+        subject { described_class.new(model, partitioning_key, retain_for: 3.months).extra_partitions }
+
+        it 'prunes the unbounded partition ending 2020-05-01' do
+          min_value_to_may = Gitlab::Database::Partitioning::TimePartition.new(model.table_name, nil, '2020-05-01',
+                                                                               partition_name: 'partitioned_test_000000')
+
+          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
+        subject { described_class.new(model, partitioning_key, retain_for: 2.months).extra_partitions }
+
+        it 'prunes the unbounded partition and the partition for May-June' do
+          expect(subject).to contain_exactly(
+            Gitlab::Database::Partitioning::TimePartition.new(model.table_name, nil, '2020-05-01', partition_name: 'partitioned_test_000000'),
+                               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
 end