Fix bulk insert for models with composite primary key

We need to work it around by passing primary key explicitly

app/models/concerns/bulk_insert_safe.rb

@@ -141,6 +141,12 @@ module BulkInsertSafe
           raise ArgumentError, "returns needs to be :ids or nil"
         end
 
+      # Handle insertions for tables with a composite primary key
+      primary_keys = connection.schema_cache.primary_keys(table_name)
+      if unique_by.blank? && primary_key != primary_keys
+        unique_by = primary_keys
+      end
+
       transaction do
         items.each_slice(batch_size).flat_map do |item_batch|
           attributes = _bulk_insert_item_attributes(

app/models/project_authorization.rb

@@ -29,6 +29,15 @@ class ProjectAuthorization < ApplicationRecord
       EOF
     end
   end
+
+  # This method overrides its ActiveRecord's version in order to work correctly
+  # with composite primary keys and fix the tests for Rails 6.1
+  #
+  # Consider using BulkInsertSafe module instead since we plan to refactor it in
+  # https://gitlab.com/gitlab-org/gitlab/-/issues/331264
+  def self.insert_all(attributes)
+    super(attributes, unique_by: connection.schema_cache.primary_keys(table_name))
+  end
 end
 
 ProjectAuthorization.prepend_mod_with('ProjectAuthorization')

spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb

@@ -580,7 +580,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
       it 'idempotently cleans up after failed background migrations' do
         expect(partitioned_model.count).to eq(0)
 
-        partitioned_model.insert!(record2.attributes)
+        partitioned_model.insert(record2.attributes, unique_by: [:id, :created_at])
 
         expect_next_instance_of(Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable) do |backfill|
           allow(backfill).to receive(:transaction_open?).and_return(false)

spec/lib/gitlab/import_export/import_failure_service_spec.rb

@@ -43,7 +43,7 @@ RSpec.describe Gitlab::ImportExport::ImportFailureService do
       let(:importable) { create(:merge_request) }
 
       it 'raise exception' do
-        expect { subject }.to raise_exception(ActiveRecord::AssociationNotFoundError, "Association named 'import_failures' was not found on MergeRequest; perhaps you misspelled it?")
+        expect { subject }.to raise_exception(ActiveRecord::AssociationNotFoundError, /Association named 'import_failures' was not found on MergeRequest/)
       end
     end
   end

spec/models/concerns/bulk_insert_safe_spec.rb

@@ -20,6 +20,13 @@ RSpec.describe BulkInsertSafe do
 
         t.index :name, unique: true
       end
+
+      create_table :bulk_insert_items_with_composite_pk, id: false, force: true do |t|
+        t.integer :id, null: true
+        t.string :name, null: true
+      end
+
+      execute("ALTER TABLE bulk_insert_items_with_composite_pk ADD PRIMARY KEY (id,name);")
     end
   end
 
@@ -27,6 +34,7 @@ RSpec.describe BulkInsertSafe do
     ActiveRecord::Schema.define do
       drop_table :bulk_insert_items, force: true
       drop_table :bulk_insert_parent_items, force: true
+      drop_table :bulk_insert_items_with_composite_pk, force: true
     end
   end
 
@@ -227,5 +235,28 @@ RSpec.describe BulkInsertSafe do
         end
       end
     end
+
+    context 'when a model with composite primary key is inserted' do
+      let_it_be(:bulk_insert_items_with_composite_pk_class) do
+        Class.new(ActiveRecord::Base) do
+          self.table_name = 'bulk_insert_items_with_composite_pk'
+
+          include BulkInsertSafe
+        end
+      end
+
+      let(:new_object) { bulk_insert_items_with_composite_pk_class.new(id: 1, name: 'composite') }
+
+      it 'successfully inserts an item' do
+        expect(ActiveRecord::InsertAll).to receive(:new)
+          .with(
+            bulk_insert_items_with_composite_pk_class, [new_object.as_json], on_duplicate: :raise, returning: false, unique_by: %w[id name]
+          ).and_call_original
+
+        expect { bulk_insert_items_with_composite_pk_class.bulk_insert!([new_object]) }.to(
+          change(bulk_insert_items_with_composite_pk_class, :count).from(0).to(1)
+        )
+      end
+    end
   end
 end