Add support for bulk-insertable associations

Given an association that is `BulkInsertSafe`, we now allow
the caller to schedule items for bulk-insertion when the
owner is `save`d.

app/models/concerns/bulk_insertable_associations.rb

+# frozen_string_literal: true
+
+##
+# ActiveRecord model classes can mix in this concern if they own associations
+# who declare themselves to be eligible for bulk-insertion via [BulkInsertSafe].
+# This allows the caller to write items from [has_many] associations en-bloc
+# when the owner is first created.
+#
+# This implementation currently has a few limitations:
+# - only works for [has_many] relations
+# - does not support the [:through] option
+# - it cannot bulk-insert items that had previously been saved, nor can the
+#   owner of the association have previously been saved; if you attempt to
+#   so, an error will be raised
+#
+# @example
+#
+#   class MergeRequestDiff < ApplicationRecord
+#     include BulkInsertableAssociations
+#
+#     # target association class must `include BulkInsertSafe`
+#     has_many :merge_request_diff_commits
+#   end
+#
+#   diff = MergeRequestDiff.new(...)
+#   diff.diff_commits << MergeRequestDiffCommit.build(...)
+#   BulkInsertableAssociations.with_bulk_insert do
+#     diff.save! # this will also write all `diff_commits` in bulk
+#   end
+#
+# Note that just like [BulkInsertSafe.bulk_insert!], validations will run for
+# all items that are scheduled for bulk-insertions.
+#
+module BulkInsertableAssociations
+  extend ActiveSupport::Concern
+
+  class << self
+    def bulk_inserts_enabled?
+      Thread.current['bulk_inserts_enabled']
+    end
+
+    # All associations that are [BulkInsertSafe] and that as a result of calls to
+    # [save] or [save!] would be written to the database, will be inserted using
+    # [bulk_insert!] instead.
+    #
+    # Note that this will only work for entities that have not been persisted yet.
+    #
+    # @param [Boolean] enabled When [true], bulk-inserts will be attempted within
+    #                          the given block. If [false], bulk-inserts will be
+    #                          disabled. This behavior can be nested.
+    def with_bulk_insert(enabled: true)
+      previous = bulk_inserts_enabled?
+      Thread.current['bulk_inserts_enabled'] = enabled
+      yield
+    ensure
+      Thread.current['bulk_inserts_enabled'] = previous
+    end
+  end
+
+  def bulk_insert_associations!
+    self.class.reflections.each do |_, reflection|
+      _bulk_insert_association!(reflection)
+    end
+  end
+
+  private
+
+  def _bulk_insert_association!(reflection)
+    return unless _association_supports_bulk_inserts?(reflection)
+
+    association = self.association(reflection.name)
+    association_items = association.target
+    return if association_items.empty?
+
+    if association_items.any?(&:persisted?)
+      raise 'Bulk-insertion of already persisted association items is not currently supported'
+    end
+
+    _bulk_insert_configure_foreign_key(reflection, association_items)
+    association.klass.bulk_insert!(association_items, validate: false)
+
+    # reset relation:
+    # 1. we successfully inserted all items
+    # 2. when accessed we force to reload the relation
+    association.reset
+  end
+
+  def _association_supports_bulk_inserts?(reflection)
+    reflection.macro == :has_many &&
+      reflection.klass < BulkInsertSafe &&
+      !reflection.through_reflection? &&
+      association_cached?(reflection.name)
+  end
+
+  def _bulk_insert_configure_foreign_key(reflection, items)
+    primary_key = self[reflection.active_record_primary_key]
+    raise "Classes including `BulkInsertableAssociations` must define a `primary_key`" unless primary_key
+
+    items.each do |item|
+      item[reflection.foreign_key] = primary_key
+
+      if reflection.type
+        item[reflection.type] = self.class.polymorphic_name
+      end
+    end
+  end
+
+  included do
+    delegate :bulk_inserts_enabled?, to: BulkInsertableAssociations
+    after_create :bulk_insert_associations!, if: :bulk_inserts_enabled?, prepend: true
+  end
+end

app/models/merge_request_diff.rb

@@ -7,6 +7,7 @@ class MergeRequestDiff < ApplicationRecord
   include EachBatch
   include Gitlab::Utils::StrongMemoize
   include ObjectStorage::BackgroundMove
+  include BulkInsertableAssociations
 
   # Don't display more than 100 commits at once
   COMMITS_SAFE_SIZE = 100

lib/gitlab/import_export/relation_tree_restorer.rb

@@ -26,8 +26,13 @@ module Gitlab
         ActiveRecord::Base.uncached do
           ActiveRecord::Base.no_touching do
             update_params!
-            update_relation_hashes!
-            create_relations!
+
+            bulk_inserts_enabled = @importable.class == ::Project &&
+              Feature.enabled?(:import_bulk_inserts, @importable.group)
+            BulkInsertableAssociations.with_bulk_insert(enabled: bulk_inserts_enabled) do
+              update_relation_hashes!
+              create_relations!
+            end
           end
         end
 
@@ -170,7 +175,7 @@ module Gitlab
 
         # if object is a hash we can create simple object
         # as it means that this is 1-to-1 vs 1-to-many
-        sub_data_hash =
+        current_item =
           if sub_data_hash.is_a?(Array)
             build_relations(
               sub_relation_key,
@@ -183,9 +188,8 @@ module Gitlab
               sub_data_hash)
           end
 
-        # persist object(s) or delete from relation
-        if sub_data_hash
-          data_hash[sub_relation_key] = sub_data_hash
+        if current_item
+          data_hash[sub_relation_key] = current_item
         else
           data_hash.delete(sub_relation_key)
         end

spec/models/concerns/bulk_insertable_associations_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe BulkInsertableAssociations do
+  class BulkFoo < ApplicationRecord
+    include BulkInsertSafe
+
+    validates :name, presence: true
+  end
+
+  class BulkBar < ApplicationRecord
+    include BulkInsertSafe
+  end
+
+  SimpleBar = Class.new(ApplicationRecord)
+
+  class BulkParent < ApplicationRecord
+    include BulkInsertableAssociations
+
+    has_many :bulk_foos
+    has_many :bulk_hunks, class_name: 'BulkFoo'
+    has_many :bulk_bars
+    has_many :simple_bars # not `BulkInsertSafe`
+    has_one :bulk_foo # not supported
+  end
+
+  before(:all) do
+    ActiveRecord::Schema.define do
+      create_table :bulk_parents, force: true do |t|
+        t.string :name, null: true
+      end
+
+      create_table :bulk_foos, force: true do |t|
+        t.string :name, null: true
+        t.belongs_to :bulk_parent, null: false
+      end
+
+      create_table :bulk_bars, force: true do |t|
+        t.string :name, null: true
+        t.belongs_to :bulk_parent, null: false
+      end
+
+      create_table :simple_bars, force: true do |t|
+        t.string :name, null: true
+        t.belongs_to :bulk_parent, null: false
+      end
+    end
+  end
+
+  after(:all) do
+    ActiveRecord::Schema.define do
+      drop_table :bulk_foos, force: true
+      drop_table :bulk_bars, force: true
+      drop_table :simple_bars, force: true
+      drop_table :bulk_parents, force: true
+    end
+  end
+
+  before do
+    ActiveRecord::Base.connection.execute('TRUNCATE bulk_foos RESTART IDENTITY')
+  end
+
+  context 'saving bulk insertable associations' do
+    let(:parent) { BulkParent.new(name: 'parent') }
+
+    context 'when items already have IDs' do
+      it 'stores nothing and raises an error' do
+        build_items(parent: parent) { |n, item| item.id = 100 + n }
+
+        expect { save_with_bulk_inserts(parent) }.to raise_error(BulkInsertSafe::PrimaryKeySetError)
+        expect(BulkFoo.count).to eq(0)
+      end
+    end
+
+    context 'when items have no IDs set' do
+      it 'stores them all and updates items with IDs' do
+        items = build_items(parent: parent)
+
+        expect(BulkFoo).to receive(:bulk_insert!).once.and_call_original
+        expect { save_with_bulk_inserts(parent) }.to change { BulkFoo.count }.from(0).to(items.size)
+        expect(parent.bulk_foos.pluck(:id)).to contain_exactly(*(1..10))
+      end
+    end
+
+    context 'when items are empty' do
+      it 'does nothing' do
+        expect(parent.bulk_foos).to be_empty
+
+        expect { save_with_bulk_inserts(parent) }.not_to change { BulkFoo.count }
+      end
+    end
+
+    context 'when relation name does not match class name' do
+      it 'stores them all' do
+        items = build_items(parent: parent, relation: :bulk_hunks)
+
+        expect(BulkFoo).to receive(:bulk_insert!).once.and_call_original
+
+        expect { save_with_bulk_inserts(parent) }.to(
+          change { BulkFoo.count }.from(0).to(items.size)
+        )
+      end
+    end
+
+    context 'with multiple threads' do
+      it 'isolates bulk insert behavior between threads' do
+        total_item_count = 10
+        parent1 = BulkParent.new(name: 'parent1')
+        parent2 = BulkParent.new(name: 'parent2')
+        build_items(parent: parent1, count: total_item_count / 2)
+        build_items(parent: parent2, count: total_item_count / 2)
+
+        expect(BulkFoo).to receive(:bulk_insert!).once.and_call_original
+        [
+          Thread.new do
+            save_with_bulk_inserts(parent1)
+          end,
+          Thread.new do
+            parent2.save!
+          end
+        ].map(&:join)
+
+        expect(BulkFoo.count).to eq(total_item_count)
+      end
+    end
+
+    context 'with multiple associations' do
+      it 'isolates writes between associations' do
+        items1 = build_items(parent: parent, relation: :bulk_foos)
+        items2 = build_items(parent: parent, relation: :bulk_bars)
+
+        expect(BulkFoo).to receive(:bulk_insert!).once.and_call_original
+        expect(BulkBar).to receive(:bulk_insert!).once.and_call_original
+
+        expect { save_with_bulk_inserts(parent) }.to(
+          change { BulkFoo.count }.from(0).to(items1.size)
+        .and(
+          change { BulkBar.count }.from(0).to(items2.size)
+        ))
+      end
+    end
+
+    context 'passing bulk insert arguments' do
+      it 'disables validations on target association' do
+        items = build_items(parent: parent)
+
+        expect(BulkFoo).to receive(:bulk_insert!).with(items, validate: false).and_return true
+
+        save_with_bulk_inserts(parent)
+      end
+    end
+
+    it 'can disable bulk-inserts within a bulk-insert block' do
+      parent1 = BulkParent.new(name: 'parent1')
+      parent2 = BulkParent.new(name: 'parent2')
+      _items1 = build_items(parent: parent1)
+      items2 = build_items(parent: parent2)
+
+      expect(BulkFoo).to receive(:bulk_insert!).once.with(items2, validate: false)
+
+      BulkInsertableAssociations.with_bulk_insert(enabled: true) do
+        BulkInsertableAssociations.with_bulk_insert(enabled: false) do
+          parent1.save!
+        end
+
+        parent2.save!
+      end
+    end
+
+    context 'when association is not bulk-insert safe' do
+      it 'saves it normally' do
+        parent.simple_bars.build
+
+        expect(SimpleBar).not_to receive(:bulk_insert!)
+        expect { save_with_bulk_inserts(parent) }.to change { SimpleBar.count }.from(0).to(1)
+      end
+    end
+
+    context 'when association is not has_many' do
+      it 'saves it normally' do
+        parent.bulk_foo = BulkFoo.new(name: 'item')
+
+        expect(BulkFoo).not_to receive(:bulk_insert!)
+        expect { save_with_bulk_inserts(parent) }.to change { BulkFoo.count }.from(0).to(1)
+      end
+    end
+
+    context 'when an item is not valid' do
+      describe '.save' do
+        it 'invalidates the parent and returns false' do
+          build_invalid_items(parent: parent)
+
+          expect(save_with_bulk_inserts(parent, bangify: false)).to be false
+          expect(parent.errors[:bulk_foos].size).to eq(1)
+
+          expect(BulkFoo.count).to eq(0)
+          expect(BulkParent.count).to eq(0)
+        end
+      end
+
+      describe '.save!' do
+        it 'invalidates the parent and raises error' do
+          build_invalid_items(parent: parent)
+
+          expect { save_with_bulk_inserts(parent) }.to raise_error(ActiveRecord::RecordInvalid)
+          expect(parent.errors[:bulk_foos].size).to eq(1)
+
+          expect(BulkFoo.count).to eq(0)
+          expect(BulkParent.count).to eq(0)
+        end
+      end
+    end
+  end
+
+  private
+
+  def save_with_bulk_inserts(entity, bangify: true)
+    BulkInsertableAssociations.with_bulk_insert { bangify ? entity.save! : entity.save }
+  end
+
+  def build_items(parent:, relation: :bulk_foos, count: 10)
+    count.times do |n|
+      item = parent.send(relation).build(name: "item_#{n}", bulk_parent_id: parent.id)
+      yield(n, item) if block_given?
+    end
+    parent.send(relation)
+  end
+
+  def build_invalid_items(parent:)
+    build_items(parent: parent).tap do |items|
+      invalid_item = items.first
+      invalid_item.name = nil
+      expect(invalid_item).not_to be_valid
+    end
+  end
+end