Drop foreign keys prior to dropping individual detached partitions

lib/gitlab/database/partitioning/detached_partition_dropper.rb

@@ -9,13 +9,10 @@ module Gitlab
           Gitlab::AppLogger.info(message: "Checking for previously detached partitions to drop")
 
           Postgresql::DetachedPartition.ready_to_drop.find_each do |detached_partition|
-            connection.transaction do
-              # Another process may have already dropped the table and deleted this entry
-              next unless (detached_partition = Postgresql::DetachedPartition.lock.find_by(id: detached_partition.id))
-
-              drop_detached_partition(detached_partition.table_name)
-
-              detached_partition.destroy!
+            if partition_attached?(qualify_partition_name(detached_partition.table_name))
+              unmark_partition(detached_partition)
+            else
+              drop_partition(detached_partition)
             end
           rescue StandardError => e
             Gitlab::AppLogger.error(message: "Failed to drop previously detached partition",
@@ -27,31 +24,100 @@ module Gitlab
 
         private
 
+        def unmark_partition(detached_partition)
+          connection.transaction do
+            # Another process may have already encountered this case and deleted this entry
+            next unless try_lock_detached_partition(detached_partition.id)
+
+            # The current partition was scheduled for deletion incorrectly
+            # Dropping it now could delete in-use data and take locks that interrupt other database activity
+            Gitlab::AppLogger.error(message: "Prevented an attempt to drop an attached database partition", partition_name: detached_partition.table_name)
+            detached_partition.destroy!
+          end
+        end
+
+        def drop_partition(detached_partition)
+          remove_foreign_keys(detached_partition)
+
+          connection.transaction do
+            # Another process may have already dropped the table and deleted this entry
+            next unless try_lock_detached_partition(detached_partition.id)
+
+            drop_detached_partition(detached_partition.table_name)
+
+            detached_partition.destroy!
+          end
+        end
+
+        def remove_foreign_keys(detached_partition)
+          partition_identifier = qualify_partition_name(detached_partition.table_name)
+
+          # We want to load all of these into memory at once to get a consistent view to loop over,
+          # since we'll be deleting from this list as we go
+          fks_to_drop = PostgresForeignKey.by_constrained_table_identifier(partition_identifier).to_a
+          fks_to_drop.each do |foreign_key|
+            drop_foreign_key_if_present(detached_partition, foreign_key)
+          end
+        end
+
+        # Drops the given foreign key for the given detached partition, but only if another process has not already
+        # detached the partition first. This method must be safe to call even if the associated partition table has already
+        # been detached, as it could be called by multiple processes at once.
+        def drop_foreign_key_if_present(detached_partition, foreign_key)
+          # It is important to only drop one foreign key per transaction.
+          # Dropping a foreign key takes an ACCESS EXCLUSIVE lock on both tables participating in the foreign key.
+
+          partition_identifier = qualify_partition_name(detached_partition.table_name)
+          with_lock_retries do
+            connection.transaction(requires_new: false) do
+              next unless try_lock_detached_partition(detached_partition.id)
+
+              # Another process may have already dropped this foreign key
+              next unless PostgresForeignKey.by_constrained_table_identifier(partition_identifier).where(name: foreign_key.name).exists?
+
+              connection.execute("ALTER TABLE #{connection.quote_table_name(partition_identifier)} DROP CONSTRAINT #{connection.quote_table_name(foreign_key.name)}")
+
+              Gitlab::AppLogger.info(message: "Dropped foreign key for previously detached partition",
+                                     partition_name: detached_partition.table_name,
+                                     referenced_table_name: foreign_key.referenced_table_identifier,
+                                     foreign_key_name: foreign_key.name)
+            end
+          end
+        end
+
         def drop_detached_partition(partition_name)
           partition_identifier = qualify_partition_name(partition_name)
 
-          if partition_detached?(partition_identifier)
           connection.drop_table(partition_identifier, if_exists: true)
 
           Gitlab::AppLogger.info(message: "Dropped previously detached partition", partition_name: partition_name)
-          else
-            Gitlab::AppLogger.error(message: "Attempt to drop attached database partition", partition_name: partition_name)
-          end
         end
 
         def qualify_partition_name(table_name)
           "#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.#{table_name}"
         end
 
-        def partition_detached?(partition_identifier)
+        def partition_attached?(partition_identifier)
           # PostgresPartition checks the pg_inherits view, so our partition will only show here if it's still attached
           # and thus should not be dropped
-          !Gitlab::Database::PostgresPartition.for_identifier(partition_identifier).exists?
+          Gitlab::Database::PostgresPartition.for_identifier(partition_identifier).exists?
+        end
+
+        def try_lock_detached_partition(id)
+          Postgresql::DetachedPartition.lock.find_by(id: id).present?
         end
 
         def connection
           Postgresql::DetachedPartition.connection
         end
+
+        def with_lock_retries(&block)
+          Gitlab::Database::WithLockRetries.new(
+            klass: self.class,
+            logger: Gitlab::AppLogger,
+            connection: connection
+          ).run(raise_on_exhaustion: true, &block)
+        end
       end
     end
   end

lib/gitlab/database/postgres_foreign_key.rb

@@ -10,6 +10,12 @@ module Gitlab
 
         where(referenced_table_identifier: identifier)
       end
+
+      scope :by_constrained_table_identifier, ->(identifier) do
+        raise ArgumentError, "Constrained table name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+
+        where(constrained_table_identifier: identifier)
+      end
     end
   end
 end

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

@@ -5,6 +5,8 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
   include Database::TableSchemaHelpers
 
+  subject(:dropper) { described_class.new }
+
   let(:connection) { ActiveRecord::Base.connection }
 
   def expect_partition_present(name)
@@ -23,10 +25,18 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
 
   before do
     connection.execute(<<~SQL)
+      CREATE TABLE referenced_table (
+        id bigserial primary key not null
+      )
+    SQL
+    connection.execute(<<~SQL)
+
       CREATE TABLE parent_table (
          id bigserial not null,
+         referenced_id bigint not null,
          created_at timestamptz not null,
-         primary key (id, created_at)
+         primary key (id, created_at),
+        constraint fk_referenced foreign key (referenced_id) references referenced_table(id)
        ) PARTITION BY RANGE(created_at)
     SQL
   end
@@ -59,7 +69,7 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
                          attached: false,
                          drop_after: 1.day.from_now)
 
-        subject.perform
+        dropper.perform
 
         expect_partition_present('test_partition')
       end
@@ -75,7 +85,7 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
       end
 
       it 'drops the partition' do
-        subject.perform
+        dropper.perform
 
         expect(table_oid('test_partition')).to be_nil
       end
@@ -86,16 +96,62 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
         end
 
         it 'does not drop the partition' do
-          subject.perform
+          dropper.perform
 
           expect(table_oid('test_partition')).not_to be_nil
         end
       end
 
+      context 'removing foreign keys' do
+        it 'removes foreign keys from the table before dropping it' do
+          expect(dropper).to receive(:drop_detached_partition).and_wrap_original do |drop_method, partition_name|
+            expect(partition_name).to eq('test_partition')
+            expect(foreign_key_exists_by_name(partition_name, 'fk_referenced', schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)).to be_falsey
+
+            drop_method.call(partition_name)
+          end
+
+          expect(foreign_key_exists_by_name('test_partition', 'fk_referenced', schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)).to be_truthy
+
+          dropper.perform
+        end
+
+        it 'does not remove foreign keys from the parent table' do
+          expect { dropper.perform }.not_to change { foreign_key_exists_by_name('parent_table', 'fk_referenced') }.from(true)
+        end
+
+        context 'when another process drops the foreign key' do
+          it 'skips dropping that foreign key' do
+            expect(dropper).to receive(:drop_foreign_key_if_present).and_wrap_original do |drop_meth, *args|
+              connection.execute('alter table gitlab_partitions_dynamic.test_partition drop constraint fk_referenced;')
+              drop_meth.call(*args)
+            end
+
+            dropper.perform
+
+            expect_partition_removed('test_partition')
+          end
+        end
+
+        context 'when another process drops the partition' do
+          it 'skips dropping the foreign key' do
+            expect(dropper).to receive(:drop_foreign_key_if_present).and_wrap_original do |drop_meth, *args|
+              connection.execute('drop table gitlab_partitions_dynamic.test_partition')
+              Postgresql::DetachedPartition.where(table_name: 'test_partition').delete_all
+            end
+
+            expect(Gitlab::AppLogger).not_to receive(:error)
+            dropper.perform
+          end
+        end
+      end
+
       context 'when another process drops the table while the first waits for a lock' do
         it 'skips the table' do
+          # First call to .lock is for removing foreign keys
+          expect(Postgresql::DetachedPartition).to receive(:lock).once.ordered.and_call_original
           # 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|
+          expect(Postgresql::DetachedPartition).to receive(:lock).once.ordered.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
@@ -106,9 +162,9 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
             locked
           end
 
-          expect(subject).not_to receive(:drop_one)
+          expect(dropper).not_to receive(:drop_one)
 
-          subject.perform
+          dropper.perform
         end
       end
     end
@@ -123,19 +179,26 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
       end
 
       it 'does not drop the partition, but does remove the DetachedPartition entry' do
-        subject.perform
+        dropper.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')
+      context 'when another process removes the entry before this process' do
+        it 'does nothing' do
+          expect(Postgresql::DetachedPartition).to receive(:lock).and_wrap_original do |lock_meth|
+            Postgresql::DetachedPartition.delete_all
+            lock_meth.call
+          end
 
-        subject.perform
+          expect(Gitlab::AppLogger).not_to receive(:error)
 
-        expect(Postgresql::DetachedPartition.exists?(id: detached_partition.id)).to be_falsey
+          dropper.perform
+
+          expect(table_oid('test_partition')).not_to be_nil
+        end
       end
     end
 
@@ -155,7 +218,7 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
       end
 
       it 'drops both partitions' do
-        subject.perform
+        dropper.perform
 
         expect_partition_removed('partition_1')
         expect_partition_removed('partition_2')
@@ -163,10 +226,10 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
 
       context 'when the first drop returns an error' do
         it 'still drops the second partition' do
-          expect(subject).to receive(:drop_detached_partition).ordered.and_raise('injected error')
-          expect(subject).to receive(:drop_detached_partition).ordered.and_call_original
+          expect(dropper).to receive(:drop_detached_partition).ordered.and_raise('injected error')
+          expect(dropper).to receive(:drop_detached_partition).ordered.and_call_original
 
-          subject.perform
+          dropper.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)

spec/lib/gitlab/database/postgres_foreign_key_spec.rb

@@ -38,4 +38,16 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model do
       expect(described_class.by_referenced_table_identifier('public.referenced_table')).to contain_exactly(expected)
     end
   end
+
+  describe '#by_constrained_table_identifier' do
+    it 'throws an error when the identifier name is not fully qualified' do
+      expect { described_class.by_constrained_table_identifier('constrained_table') }.to raise_error(ArgumentError, /not fully qualified/)
+    end
+
+    it 'finds the foreign keys for the constrained table' do
+      expected = described_class.where(name: %w[fk_constrained_to_referenced fk_constrained_to_other_referenced]).to_a
+
+      expect(described_class.by_constrained_table_identifier('public.constrained_table')).to match_array(expected)
+    end
+  end
 end