Partition manager identifies partitions to detach

Enumerate the partitions that are outside of the retention period, and
for each one log a message that we are planning to detach it.

This way we can run this code in production for a while
and be sure that the choices of partitions to detach are correct.

app/models/concerns/partitioned_table.rb

@@ -15,7 +15,7 @@ module PartitionedTable
 
       @partitioning_strategy = strategy_class.new(self, partitioning_key, **kwargs)
 
-      Gitlab::Database::Partitioning::PartitionCreator.register(self)
+      Gitlab::Database::Partitioning::PartitionManager.register(self)
     end
   end
 end

app/workers/partition_creation_worker.rb

@@ -10,7 +10,7 @@ class PartitionCreationWorker
   idempotent!
 
   def perform
-    Gitlab::Database::Partitioning::PartitionCreator.new.create_partitions
+    Gitlab::Database::Partitioning::PartitionManager.new.create_partitions
   ensure
     Gitlab::Database::Partitioning::PartitionMonitoring.new.report_metrics
   end

config/initializers/postgres_partitioning.rb

@@ -3,15 +3,15 @@
 # Make sure we have loaded partitioned models here
 # (even with eager loading disabled).
 
-Gitlab::Database::Partitioning::PartitionCreator.register(AuditEvent)
-Gitlab::Database::Partitioning::PartitionCreator.register(WebHookLog)
+Gitlab::Database::Partitioning::PartitionManager.register(AuditEvent)
+Gitlab::Database::Partitioning::PartitionManager.register(WebHookLog)
 
 if Gitlab.ee?
-  Gitlab::Database::Partitioning::PartitionCreator.register(IncidentManagement::PendingEscalations::Alert)
+  Gitlab::Database::Partitioning::PartitionManager.register(IncidentManagement::PendingEscalations::Alert)
 end
 
 begin
-  Gitlab::Database::Partitioning::PartitionCreator.new.create_partitions unless ENV['DISABLE_POSTGRES_PARTITION_CREATION_ON_STARTUP']
+  Gitlab::Database::Partitioning::PartitionManager.new.create_partitions unless ENV['DISABLE_POSTGRES_PARTITION_CREATION_ON_STARTUP']
 rescue ActiveRecord::ActiveRecordError, PG::Error
   # ignore - happens when Rake tasks yet have to create a database, e.g. for testing
 end

lib/gitlab/database/partitioning/partition_creator.rb → lib/gitlab/database/partitioning/partition_manager.rb

@@ -3,7 +3,7 @@
 module Gitlab
   module Database
     module Partitioning
-      class PartitionCreator
+      class PartitionManager
         def self.register(model)
           raise ArgumentError, "Only models with a #partitioning_strategy can be registered." unless model.respond_to?(:partitioning_strategy)
 
@@ -15,7 +15,8 @@ module Gitlab
         end
 
         LEASE_TIMEOUT = 1.minute
-        LEASE_KEY = 'database_partition_creation_%s'
+        CREATION_LEASE_KEY = 'database_partition_creation_%s'
+        DETACH_LEASE_KEY = 'database_partition_detach_%s'
 
         attr_reader :models
 
@@ -31,7 +32,7 @@ module Gitlab
             # The prevailing situation is no missing partitions
             next if missing_partitions(model).empty?
 
-            only_with_exclusive_lease(model) do
+            only_with_exclusive_creation_lease(model) do
               partitions_to_create = missing_partitions(model)
 
               next if partitions_to_create.empty?
@@ -43,6 +44,22 @@ module Gitlab
           end
         end
 
+        def detach_partitions
+          models.each do |model|
+            next if extra_partitions(model).empty?
+
+            only_with_exclusive_removal_lease(model) do
+              partitions_to_detach = extra_partitions(model)
+
+              next if partitions_to_detach.empty?
+
+              detach(partitions_to_detach)
+            end
+          rescue StandardError => e
+            Gitlab::AppLogger.error("Failed to remove partition(s) for #{model.table_name}: #{e.class}: #{e.message}")
+          end
+        end
+
         private
 
         def missing_partitions(model)
@@ -51,14 +68,32 @@ module Gitlab
           model.partitioning_strategy.missing_partitions
         end
 
-        def only_with_exclusive_lease(model)
-          lease = Gitlab::ExclusiveLease.new(LEASE_KEY % model.table_name, timeout: LEASE_TIMEOUT)
+        def extra_partitions(model)
+          return [] unless connection.table_exists?(model.table_name)
+
+          model.partitioning_strategy.extra_partitions
+        end
+
+        def only_with_exclusive_lease(model, lease_key:)
+          lease = Gitlab::ExclusiveLease.new(lease_key % model.table_name, timeout: LEASE_TIMEOUT)
 
           yield if lease.try_obtain
         ensure
           lease&.cancel
         end
 
+        def only_with_exclusive_creation_lease(model)
+          only_with_exclusive_lease(model, lease_key: CREATION_LEASE_KEY) do
+            yield
+          end
+        end
+
+        def only_with_exclusive_removal_lease(model)
+          only_with_exclusive_lease(model, lease_key: DETACH_LEASE_KEY) do
+            yield
+          end
+        end
+
         def create(model, partitions)
           connection.transaction do
             with_lock_retries do
@@ -71,6 +106,18 @@ module Gitlab
           end
         end
 
+        def detach(partitions)
+          connection.transaction do
+            with_lock_retries do
+              partitions.each { |p| detach_one_partition(p) }
+            end
+          end
+        end
+
+        def detach_one_partition(partition)
+          Gitlab::AppLogger.info("Planning to detach #{partition.partition_name} for table #{partition.table}")
+        end
+
         def with_lock_retries(&block)
           Gitlab::Database::WithLockRetries.new(
             klass: self.class,

lib/gitlab/database/partitioning/partition_monitoring.rb

@@ -6,7 +6,7 @@ module Gitlab
       class PartitionMonitoring
         attr_reader :models
 
-        def initialize(models = PartitionCreator.models)
+        def initialize(models = PartitionManager.models)
           @models = models
         end
 

lib/tasks/gitlab/db.rake

@@ -118,7 +118,7 @@ namespace :gitlab do
 
     desc 'Create missing dynamic database partitions'
     task create_dynamic_partitions: :environment do
-      Gitlab::Database::Partitioning::PartitionCreator.new.create_partitions
+      Gitlab::Database::Partitioning::PartitionManager.new.create_partitions
     end
 
     # This is targeted towards deploys and upgrades of GitLab.

spec/lib/gitlab/database/partitioning/partition_creator_spec.rb → spec/lib/gitlab/database/partitioning/partition_manager_spec.rb

@@ -2,8 +2,9 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::Partitioning::PartitionCreator do
+RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
   include Database::PartitioningHelpers
+  include Database::TableSchemaHelpers
   include ExclusiveLeaseHelpers
 
   describe '.register' do
@@ -27,7 +28,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionCreator do
       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::LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
+      stub_exclusive_lease(described_class::CREATION_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
 
     let(:partitions) do
@@ -93,4 +94,52 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionCreator do
       subject
     end
   end
+
+  describe '#detach_partitions (mocked)' do
+    subject { manager.detach_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) }
+    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::DETACH_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
+    end
+
+    let(:extra_partitions) do
+      [
+        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
+
+    it 'detaches each extra partition' do
+      extra_partitions.each { |p| expect(manager).to receive(:detach_one_partition).with(p) }
+
+      subject
+    end
+
+    context 'error handling' do
+      let(:models) do
+        [
+          double(partitioning_strategy: error_strategy, table_name: table),
+          model
+        ]
+      end
+
+      let(:error_strategy) { double }
+
+      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) }
+
+        subject
+      end
+    end
+  end
 end

spec/models/concerns/partitioned_table_spec.rb

@@ -37,7 +37,7 @@ RSpec.describe PartitionedTable do
     end
 
     it 'registers itself with the PartitionCreator' do
-      expect(Gitlab::Database::Partitioning::PartitionCreator).to receive(:register).with(my_class)
+      expect(Gitlab::Database::Partitioning::PartitionManager).to receive(:register).with(my_class)
 
       subject
     end

spec/workers/partition_creation_worker_spec.rb

@@ -6,16 +6,16 @@ RSpec.describe PartitionCreationWorker do
   describe '#perform' do
     subject { described_class.new.perform }
 
-    let(:creator) {  instance_double('PartitionCreator', create_partitions: nil) }
+    let(:manager) {  instance_double('PartitionManager', create_partitions: nil) }
     let(:monitoring) { instance_double('PartitionMonitoring', report_metrics: nil) }
 
     before do
-      allow(Gitlab::Database::Partitioning::PartitionCreator).to receive(:new).and_return(creator)
+      allow(Gitlab::Database::Partitioning::PartitionManager).to receive(:new).and_return(manager)
       allow(Gitlab::Database::Partitioning::PartitionMonitoring).to receive(:new).and_return(monitoring)
     end
 
     it 'delegates to PartitionCreator' do
-      expect(creator).to receive(:create_partitions)
+      expect(manager).to receive(:create_partitions)
 
       subject
     end