Retain non-empty partitions

This change adds an option `retain_non_empty_partitions` to the montly
partitioning strategy.

lib/gitlab/database/partitioning/monthly_strategy.rb

@@ -4,17 +4,18 @@ module Gitlab
   module Database
     module Partitioning
       class MonthlyStrategy
-        attr_reader :model, :partitioning_key, :retain_for
+        attr_reader :model, :partitioning_key, :retain_for, :retain_non_empty_partitions
 
         # We create this many partitions in the future
         HEADROOM = 6.months
 
         delegate :table_name, to: :model
 
-        def initialize(model, partitioning_key, retain_for: nil)
+        def initialize(model, partitioning_key, retain_for: nil, retain_non_empty_partitions: false)
           @model = model
           @partitioning_key = partitioning_key
           @retain_for = retain_for
+          @retain_non_empty_partitions = retain_non_empty_partitions
         end
 
         def current_partitions
@@ -29,7 +30,10 @@ module Gitlab
         end
 
         def extra_partitions
-          current_partitions - desired_partitions
+          partitions = current_partitions - desired_partitions
+          partitions.reject!(&:holds_data?) if retain_non_empty_partitions
+
+          partitions
         end
 
         private

lib/gitlab/database/partitioning/time_partition.rb

@@ -69,6 +69,10 @@ module Gitlab
           partition_name <=> other.partition_name
         end
 
+        def holds_data?
+          conn.execute("SELECT 1 FROM #{fully_qualified_partition} LIMIT 1").ntuples > 0
+        end
+
         private
 
         def date_or_nil(obj)

spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
+  let(:connection) { ActiveRecord::Base.connection }
+
   describe '#current_partitions' do
     subject { described_class.new(model, partitioning_key).current_partitions }
 
@@ -11,7 +13,7 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
     let(:table_name) { :partitioned_test }
 
     before do
-      ActiveRecord::Base.connection.execute(<<~SQL)
+      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);
@@ -52,7 +54,7 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
 
     context 'with existing partitions' do
       before do
-        ActiveRecord::Base.connection.execute(<<~SQL)
+        connection.execute(<<~SQL)
           CREATE TABLE #{model.table_name}
             (id serial not null, created_at timestamptz not null, PRIMARY KEY (id, created_at))
             PARTITION BY RANGE (created_at);
@@ -113,7 +115,7 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
 
     context 'without existing partitions' do
       before do
-        ActiveRecord::Base.connection.execute(<<~SQL)
+        connection.execute(<<~SQL)
           CREATE TABLE #{model.table_name}
             (id serial not null, created_at timestamptz not null, PRIMARY KEY (id, created_at))
             PARTITION BY RANGE (created_at);
@@ -159,7 +161,7 @@ RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
 
     context 'with a regular partition but no catchall (MINVALUE, to) partition' do
       before do
-        ActiveRecord::Base.connection.execute(<<~SQL)
+        connection.execute(<<~SQL)
           CREATE TABLE #{model.table_name}
             (id serial not null, created_at timestamptz not null, PRIMARY KEY (id, created_at))
             PARTITION BY RANGE (created_at);
@@ -248,6 +250,25 @@ 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 retain_non_empty_partitions is true' do
+          subject { described_class.new(model, partitioning_key, retain_for: 2.months, retain_non_empty_partitions: true).extra_partitions }
+
+          it 'prunes empty partitions' 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
+
+          it 'does not prune non-empty partitions' do
+            connection.execute("INSERT INTO #{table_name} (created_at) VALUES (('2020-05-15'))") # inserting one record into partitioned_test_202005
+
+            expect(subject).to contain_exactly(
+              Gitlab::Database::Partitioning::TimePartition.new(model.table_name, nil, '2020-05-01', partition_name: 'partitioned_test_000000')
+            )
+          end
+        end
       end
     end
   end