DB structure for loose foreign keys

app/models/loose_foreign_keys.rb

+# frozen_string_literal: true
+
+module LooseForeignKeys
+  def self.table_name_prefix
+    'loose_foreign_keys_'
+  end
+end

app/models/loose_foreign_keys/deleted_record.rb

+# frozen_string_literal: true
+
+class LooseForeignKeys::DeletedRecord < ApplicationRecord
+  extend SuppressCompositePrimaryKeyWarning
+  include PartitionedTable
+
+  partitioned_by :created_at, strategy: :monthly, retain_for: 3.months, retain_non_empty_partitions: true
+
+  scope :ordered_by_primary_keys, -> { order(:created_at, :deleted_table_name, :deleted_table_primary_key_value) }
+
+  def self.load_batch(batch_size)
+    ordered_by_primary_keys
+      .limit(batch_size)
+      .to_a
+  end
+
+  # Because the table has composite primary keys, the delete_all or delete methods are not going to work.
+  # This method implements deletion that benefits from the primary key index, example:
+  #
+  # > DELETE
+  # > FROM "loose_foreign_keys_deleted_records"
+  # > WHERE (created_at,
+  # >        deleted_table_name,
+  # >        deleted_table_primary_key_value) IN
+  # >     (SELECT created_at::TIMESTAMP WITH TIME ZONE,
+  # >                                             deleted_table_name,
+  # >                                             deleted_table_primary_key_value
+  # >      FROM (VALUES (LIST_OF_VALUES)) AS primary_key_values (created_at, deleted_table_name, deleted_table_primary_key_value))
+  def self.delete_records(records)
+    values = records.pluck(:created_at, :deleted_table_name, :deleted_table_primary_key_value)
+
+    primary_keys = connection.primary_keys(table_name).join(', ')
+
+    primary_keys_with_type_cast = [
+      Arel.sql('created_at::timestamp with time zone'),
+      Arel.sql('deleted_table_name'),
+      Arel.sql('deleted_table_primary_key_value')
+    ]
+
+    value_list = Arel::Nodes::ValuesList.new(values)
+
+    # (SELECT primary keys FROM VALUES)
+    inner_query = Arel::SelectManager.new
+    inner_query.from("#{Arel::Nodes::Grouping.new([value_list]).as('primary_key_values').to_sql} (#{primary_keys})")
+    inner_query.projections = primary_keys_with_type_cast
+
+    where(Arel::Nodes::Grouping.new([Arel.sql(primary_keys)]).in(inner_query)).delete_all
+  end
+end

config/initializers/postgres_partitioning.rb

@@ -5,6 +5,7 @@
 
 Gitlab::Database::Partitioning::PartitionManager.register(AuditEvent)
 Gitlab::Database::Partitioning::PartitionManager.register(WebHookLog)
+Gitlab::Database::Partitioning::PartitionManager.register(LooseForeignKeys::DeletedRecord)
 
 if Gitlab.ee?
   Gitlab::Database::Partitioning::PartitionManager.register(IncidentManagement::PendingEscalations::Alert)

db/migrate/20210826122748_create_loose_foreign_keys_deleted_records.rb

+# frozen_string_literal: true
+
+class CreateLooseForeignKeysDeletedRecords < ActiveRecord::Migration[6.1]
+  include Gitlab::Database::PartitioningMigrationHelpers::TableManagementHelpers
+
+  def up
+    constraint_name = check_constraint_name('loose_foreign_keys_deleted_records', 'deleted_table_name', 'max_length')
+    execute(<<~SQL)
+      CREATE TABLE loose_foreign_keys_deleted_records (
+        created_at timestamp with time zone NOT NULL DEFAULT NOW(),
+        deleted_table_name text NOT NULL,
+        deleted_table_primary_key_value bigint NOT NULL,
+        PRIMARY KEY (created_at, deleted_table_name, deleted_table_primary_key_value),
+        CONSTRAINT #{constraint_name} CHECK ((char_length(deleted_table_name) <= 63))
+      ) PARTITION BY RANGE (created_at);
+    SQL
+
+    min_date = Date.today - 1.month
+    max_date = Date.today + 3.months
+    create_daterange_partitions('loose_foreign_keys_deleted_records', 'created_at', min_date, max_date)
+  end
+
+  def down
+    drop_table :loose_foreign_keys_deleted_records
+  end
+end

db/migrate/20210826145509_add_function_for_inserting_deleted_records.rb

+# frozen_string_literal: true
+
+class AddFunctionForInsertingDeletedRecords < ActiveRecord::Migration[6.1]
+  include Gitlab::Database::MigrationHelpers
+  include Gitlab::Database::MigrationHelpers::LooseForeignKeyHelpers
+
+  def up
+    execute(<<~SQL)
+      CREATE FUNCTION #{DELETED_RECORDS_INSERT_FUNCTION_NAME}()
+      RETURNS TRIGGER AS
+      $$
+      BEGIN
+        INSERT INTO loose_foreign_keys_deleted_records
+        (deleted_table_name, deleted_table_primary_key_value)
+        SELECT TG_TABLE_NAME, old_table.id FROM old_table
+        ON CONFLICT DO NOTHING;
+
+        RETURN NULL;
+      END
+      $$ LANGUAGE PLPGSQL
+    SQL
+  end
+
+  def down
+    drop_function(DELETED_RECORDS_INSERT_FUNCTION_NAME)
+  end
+end

db/schema_migrations/20210826122748

+a1290cc671c487a7c24bfdb02c564d656a6606258e680e65ed108e3a28de10ca
\ No newline at end of file

db/schema_migrations/20210826145509

+661b2f03f2387f0d49cbb11c333ad29c6af5caed1f43e860fa0f263f8e7371c2
\ No newline at end of file

db/structure.sql

@@ -22,6 +22,19 @@ RETURN NULL;
 END
 $$;
 
+CREATE FUNCTION insert_into_loose_foreign_keys_deleted_records() RETURNS trigger
+    LANGUAGE plpgsql
+    AS $$
+BEGIN
+  INSERT INTO loose_foreign_keys_deleted_records
+  (deleted_table_name, deleted_table_primary_key_value)
+  SELECT TG_TABLE_NAME, old_table.id FROM old_table
+  ON CONFLICT DO NOTHING;
+
+  RETURN NULL;
+END
+$$;
+
 CREATE FUNCTION integrations_set_type_new() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
@@ -162,6 +175,14 @@ CREATE TABLE incident_management_pending_issue_escalations (
 )
 PARTITION BY RANGE (process_at);
 
+CREATE TABLE loose_foreign_keys_deleted_records (
+    created_at timestamp with time zone DEFAULT now() NOT NULL,
+    deleted_table_name text NOT NULL,
+    deleted_table_primary_key_value bigint NOT NULL,
+    CONSTRAINT check_7229f9527e CHECK ((char_length(deleted_table_name) <= 63))
+)
+PARTITION BY RANGE (created_at);
+
 CREATE TABLE web_hook_logs (
     id bigint NOT NULL,
     web_hook_id integer NOT NULL,
@@ -23067,6 +23088,9 @@ ALTER TABLE ONLY list_user_preferences
 ALTER TABLE ONLY lists
     ADD CONSTRAINT lists_pkey PRIMARY KEY (id);
 
+ALTER TABLE ONLY loose_foreign_keys_deleted_records
+    ADD CONSTRAINT loose_foreign_keys_deleted_records_pkey PRIMARY KEY (created_at, deleted_table_name, deleted_table_primary_key_value);
+
 ALTER TABLE ONLY members
     ADD CONSTRAINT members_pkey PRIMARY KEY (id);
 

lib/gitlab/database/migration_helpers/loose_foreign_key_helpers.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module MigrationHelpers
+      module LooseForeignKeyHelpers
+        include Gitlab::Database::SchemaHelpers
+
+        DELETED_RECORDS_INSERT_FUNCTION_NAME = 'insert_into_loose_foreign_keys_deleted_records'
+
+        def track_record_deletions(table)
+          execute(<<~SQL)
+          CREATE TRIGGER #{record_deletion_trigger_name(table)}
+          AFTER DELETE ON #{table} REFERENCING OLD TABLE AS old_table
+          FOR EACH STATEMENT
+          EXECUTE FUNCTION #{DELETED_RECORDS_INSERT_FUNCTION_NAME}();
+          SQL
+        end
+
+        def untrack_record_deletions(table)
+          drop_trigger(table, record_deletion_trigger_name(table))
+        end
+
+        private
+
+        def record_deletion_trigger_name(table)
+          "#{table}_loose_fk_trigger"
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/migration_helpers/loose_foreign_key_helpers_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::MigrationHelpers::LooseForeignKeyHelpers do
+  let_it_be(:migration) do
+    ActiveRecord::Migration.new.extend(described_class)
+  end
+
+  let(:model) do
+    Class.new(ApplicationRecord) do
+      self.table_name = 'loose_fk_test_table'
+    end
+  end
+
+  before(:all) do
+    migration.create_table :loose_fk_test_table do |t|
+      t.timestamps
+    end
+  end
+
+  before do
+    3.times { model.create! }
+  end
+
+  context 'when the record deletion tracker trigger is not installed' do
+    it 'does store record deletions' do
+      model.delete_all
+
+      expect(LooseForeignKeys::DeletedRecord.count).to eq(0)
+    end
+  end
+
+  context 'when the record deletion tracker trigger is installed' do
+    before do
+      migration.track_record_deletions(:loose_fk_test_table)
+    end
+
+    it 'stores the record deletion' do
+      records = model.all
+      record_to_be_deleted = records.last
+
+      record_to_be_deleted.delete
+
+      expect(LooseForeignKeys::DeletedRecord.count).to eq(1)
+      deleted_record = LooseForeignKeys::DeletedRecord.all.first
+
+      expect(deleted_record.deleted_table_primary_key_value).to eq(record_to_be_deleted.id)
+      expect(deleted_record.deleted_table_name).to eq('loose_fk_test_table')
+    end
+
+    it 'stores multiple record deletions' do
+      model.delete_all
+
+      expect(LooseForeignKeys::DeletedRecord.count).to eq(3)
+    end
+  end
+end

spec/models/loose_foreign_keys/deleted_record_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe LooseForeignKeys::DeletedRecord do
+  let_it_be(:deleted_record_1) { described_class.create!(created_at: 1.day.ago, deleted_table_name: 'projects', deleted_table_primary_key_value: 5) }
+  let_it_be(:deleted_record_2) { described_class.create!(created_at: 3.days.ago, deleted_table_name: 'projects', deleted_table_primary_key_value: 1) }
+  let_it_be(:deleted_record_3) { described_class.create!(created_at: 5.days.ago, deleted_table_name: 'projects', deleted_table_primary_key_value: 3) }
+  let_it_be(:deleted_record_4) { described_class.create!(created_at: 10.days.ago, deleted_table_name: 'projects', deleted_table_primary_key_value: 1) } # duplicate
+
+  # skip created_at because it gets truncated after insert
+  def map_attributes(records)
+    records.pluck(:deleted_table_name, :deleted_table_primary_key_value)
+  end
+
+  describe 'partitioning strategy' do
+    it 'has retain_non_empty_partitions option' do
+      expect(described_class.partitioning_strategy.retain_non_empty_partitions).to eq(true)
+    end
+  end
+
+  describe '.load_batch' do
+    it 'loads records and orders them by creation date' do
+      records = described_class.load_batch(4)
+
+      expect(map_attributes(records)).to eq([['projects', 1], ['projects', 3], ['projects', 1], ['projects', 5]])
+    end
+
+    it 'supports configurable batch size' do
+      records = described_class.load_batch(2)
+
+      expect(map_attributes(records)).to eq([['projects', 1], ['projects', 3]])
+    end
+  end
+
+  describe '.delete_records' do
+    it 'deletes exactly one record' do
+      described_class.delete_records([deleted_record_2])
+
+      expect(described_class.count).to eq(3)
+      expect(described_class.find_by(created_at: deleted_record_2.created_at)).to eq(nil)
+    end
+
+    it 'deletes two records' do
+      described_class.delete_records([deleted_record_2, deleted_record_4])
+
+      expect(described_class.count).to eq(2)
+    end
+
+    it 'deletes all records' do
+      described_class.delete_records([deleted_record_1, deleted_record_2, deleted_record_3, deleted_record_4])
+
+      expect(described_class.count).to eq(0)
+    end
+  end
+end