BulkInsertSafe mixin guards against bulk-unsafe calls

This mixin overrides certain active record calls that modify the
callback chain in ways incompatible with bulk inserts

app/models/concerns/bulk_insert_safe.rb

+# frozen_string_literal: true
+
+module BulkInsertSafe
+  extend ActiveSupport::Concern
+
+  # These are the callbacks we think safe when used on models that are
+  # written to the database in bulk
+  CALLBACK_NAME_WHITELIST = Set[
+    :initialize,
+    :validate,
+    :validation,
+    :find,
+    :destroy
+  ].freeze
+
+  MethodNotAllowedError = Class.new(StandardError)
+
+  class_methods do
+    def set_callback(name, *args)
+      unless _bulk_insert_callback_allowed?(name, args)
+        raise MethodNotAllowedError.new(
+          "Not allowed to call `set_callback(#{name}, #{args})` when model extends `BulkInsertSafe`." \
+            "Callbacks that fire per each record being inserted do not work with bulk-inserts.")
+      end
+
+      super
+    end
+
+    private
+
+    def _bulk_insert_callback_allowed?(name, args)
+      _bulk_insert_whitelisted?(name) || _bulk_insert_saved_from_belongs_to?(name, args)
+    end
+
+    # belongs_to associations will install a before_save hook during class loading
+    def _bulk_insert_saved_from_belongs_to?(name, args)
+      args.first == :before && args.second.to_s.start_with?('autosave_associated_records_for_')
+    end
+
+    def _bulk_insert_whitelisted?(name)
+      CALLBACK_NAME_WHITELIST.include?(name)
+    end
+  end
+end

app/models/label_link.rb

 # frozen_string_literal: true
 
 class LabelLink < ApplicationRecord
+  include BulkInsertSafe
   include Importable
 
   belongs_to :target, polymorphic: true, inverse_of: :label_links # rubocop:disable Cop/PolymorphicAssociations

app/models/merge_request_diff_commit.rb

 # frozen_string_literal: true
 
 class MergeRequestDiffCommit < ApplicationRecord
+  include BulkInsertSafe
   include ShaAttribute
   include CachedCommit
 

app/models/merge_request_diff_file.rb

 # frozen_string_literal: true
 
 class MergeRequestDiffFile < ApplicationRecord
+  include BulkInsertSafe
   include Gitlab::EncodingHelper
   include DiffFile
 

spec/models/concerns/bulk_insert_safe_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe BulkInsertSafe do
+  class BulkInsertItem < ApplicationRecord
+    include BulkInsertSafe
+  end
+
+  module InheritedUnsafeMethods
+    extend ActiveSupport::Concern
+
+    included do
+      after_save -> { "unsafe" }
+    end
+  end
+
+  module InheritedSafeMethods
+    extend ActiveSupport::Concern
+
+    included do
+      after_initialize -> { "safe" }
+    end
+  end
+
+  it_behaves_like 'a BulkInsertSafe model', BulkInsertItem
+
+  context 'when inheriting class methods' do
+    it 'raises an error when method is not bulk-insert safe' do
+      expect { BulkInsertItem.include(InheritedUnsafeMethods) }.to(
+        raise_error(subject::MethodNotAllowedError))
+    end
+
+    it 'does not raise an error when method is bulk-insert safe' do
+      expect { BulkInsertItem.include(InheritedSafeMethods) }.not_to raise_error
+    end
+  end
+end

spec/models/label_link_spec.rb

@@ -7,4 +7,6 @@ describe LabelLink do
 
   it { is_expected.to belong_to(:label) }
   it { is_expected.to belong_to(:target) }
+
+  it_behaves_like 'a BulkInsertSafe model', LabelLink
 end

spec/models/merge_request_diff_commit_spec.rb

@@ -6,6 +6,8 @@ describe MergeRequestDiffCommit do
   let(:merge_request) { create(:merge_request) }
   let(:project) { merge_request.project }
 
+  it_behaves_like 'a BulkInsertSafe model', MergeRequestDiffCommit
+
   describe '#to_hash' do
     subject { merge_request.commits.first }
 

spec/models/merge_request_diff_file_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 describe MergeRequestDiffFile do
+  it_behaves_like 'a BulkInsertSafe model', MergeRequestDiffFile
+
   describe '#diff' do
     context 'when diff is not stored' do
       let(:unpacked) { 'unpacked' }

spec/support/shared_examples/models/concerns/bulk_insert_safe_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'a BulkInsertSafe model' do |target_class|
+  # We consider all callbacks unsafe for bulk insertions unless we have explicitly
+  # whitelisted them (esp. anything related to :save, :create, :commit etc.)
+  let(:callback_method_blacklist) do
+    ActiveRecord::Callbacks::CALLBACKS.reject do |callback|
+      cb_name = callback.to_s.gsub(/(before_|after_|around_)/, '').to_sym
+      BulkInsertSafe::CALLBACK_NAME_WHITELIST.include?(cb_name)
+    end.to_set
+  end
+
+  context 'when calling class methods directly' do
+    it 'raises an error when method is not bulk-insert safe' do
+      callback_method_blacklist.each do |m|
+        expect { target_class.send(m, nil) }.to(
+          raise_error(BulkInsertSafe::MethodNotAllowedError),
+          "Expected call to #{m} to raise an error, but it didn't"
+        )
+      end
+    end
+
+    it 'does not raise an error when method is bulk-insert safe' do
+      BulkInsertSafe::CALLBACK_NAME_WHITELIST.each do |name|
+        expect { target_class.set_callback(name) {} }.not_to raise_error
+      end
+    end
+
+    it 'does not raise an error when the call is triggered by belongs_to' do
+      expect { target_class.belongs_to(:other_record) }.not_to raise_error
+    end
+  end
+end