Add locking and test case

db/post_migrate/20220110233155_remove_dast_site_profiles_builds_ci_build_id_fk.rb

@@ -7,6 +7,7 @@ class RemoveDastSiteProfilesBuildsCiBuildIdFk < Gitlab::Database::Migration[1.0]
 
   def up
     with_lock_retries do
+      execute('LOCK ci_builds, dast_site_profiles_builds IN ACCESS EXCLUSIVE MODE')
       remove_foreign_key_if_exists(:dast_site_profiles_builds, :ci_builds, name: CONSTRAINT_NAME)
     end
   end

ee/spec/models/dast/site_profiles_build_spec.rb

@@ -27,4 +27,9 @@ RSpec.describe Dast::SiteProfilesBuild, type: :model do
       end
     end
   end
+
+  it_behaves_like 'cleanup by a loose foreign key' do
+    let!(:model) { create(:dast_site_profiles_build) }
+    let(:parent) { model.ci_build }
+  end
 end

spec/support/shared_examples/loose_foreign_keys/have_loose_foreign_key.rb

@@ -55,8 +55,13 @@ RSpec.shared_examples 'cleanup by a loose foreign key' do
   end
 
   def find_model
-    primary_key = model.class.primary_key.to_sym
-    model.class.find_by(primary_key => model.public_send(primary_key))
+    query = model.class
+    # handle composite primary keys
+    connection = model.class.connection
+    connection.primary_keys(model.class.table_name).each do |primary_key|
+      query = query.where(primary_key => model.public_send(primary_key))
+    end
+    query.first
   end
 
   it 'cleans up (delete or nullify) the model' do