refactor relation factory

lib/gitlab/import_export/group_project_finder.rb

 module Gitlab
   module ImportExport
     class GroupProjectFinder
-      def self.find(*args)
-        new(*args).find
+      def self.find_or_new(*args)
+        new(*args).find_or_new
+      end
+
+      def self.find_or_create(*args)
+        new(*args).find_or_create
       end
 
       def initialize(klass, attributes)
@@ -12,22 +16,41 @@ module Gitlab
         @project_id = @attributes['project_id']
       end
 
-      def find
-        @klass.where(where_clause)
+      def find_or_new
+        @klass.where(where_clause).first || @klass.new(project_attributes)
+      end
+
+      def find_or_create
+        @klass.where(where_clause).first || @klass.create(project_attributes)
       end
 
       private
 
       def where_clause
         @attributes.except('group_id', 'project_id').map do |key, value|
-          table[key].eq(value).and(table[:group_id].eq(@group_id))
-            .or(table[key].eq(value).and(table[:project_id].eq(@project_id)))
+          project_clause = table[key].eq(value).and(table[:project_id].eq(@project_id))
+
+          if @group_id
+              project_clause.or(table[key].eq(value).and(table[:group_id].eq(@group_id)))
+          else
+            project_clause
+          end
         end.reduce(:or)
       end
 
       def table
         @table ||= @klass.arel_table
       end
+
+      def project_attributes
+        @attributes.except('group_id').tap do |atts|
+          atts['type'] = 'ProjectLabel' if label?
+        end
+      end
+
+      def label?
+        @klass == Label || @klass < Label
+      end
     end
   end
 end

lib/gitlab/import_export/relation_factory.rb

@@ -80,8 +80,8 @@ module Gitlab
         case @relation_name
         when :merge_request_diff_files       then setup_diff
         when :notes                          then setup_note
-        when :project_label, :project_labels then setup_label
-        when :milestone, :milestones         then setup_milestone
+        when :milestone, :milestones,
+          :project_label, :project_labels then setup_project_group
         when 'Ci::Pipeline'                  then setup_pipeline
         else
           @relation_hash['project_id'] = @project.id
@@ -167,23 +167,10 @@ module Gitlab
         @relation_hash['target_project_id'] && @relation_hash['target_project_id'] == @relation_hash['source_project_id']
       end
 
-      def setup_label
-        # If there's no group, move the label to a project label
-        if @relation_hash['type'] == 'GroupLabel' && @relation_hash['group_id']
-          @relation_hash['project_id'] = nil
-          @relation_name = :group_label
-        else
-          @relation_hash['group_id'] = nil
-          @relation_hash['type'] = 'ProjectLabel'
-        end
-      end
-
-      def setup_milestone
-        if @relation_hash['group_id']
-          @relation_hash['group_id'] = @project.group.id
-        else
+      def setup_project_group
         @relation_hash['project_id'] = @project.id
-        end
+        @relation_hash['group_id'] = @project.group&.id
+        @relation_hash.delete('type')
       end
 
       def reset_tokens!
@@ -288,30 +275,29 @@ module Gitlab
       end
 
       def find_or_create_object!
-        finder_attributes = if @relation_name == :group_label
-                              %w[title group_id]
-                            elsif parsed_relation_hash['project_id']
-                              %w[title project_id]
-                            else
-                              %w[title group_id]
-                            end
-
-        finder_hash = parsed_relation_hash.slice(*finder_attributes)
+        finder_hash = parsed_relation_hash.slice('title', 'project_id', 'group_id')
 
         if label?
-          label = relation_class.find_or_initialize_by(finder_hash)
+          label = GroupProjectFinder.find_or_new(Label, finder_hash)
           parsed_relation_hash.delete('priorities') if label.persisted?
+          parsed_relation_hash.delete('type')
 
           label.save!
           label
         else
-          relation_class.find_or_create_by(finder_hash)
+          parsed_relation_hash.delete('type') if milestone?
+
+          GroupProjectFinder.find_or_create(relation_class, finder_hash)
         end
       end
 
       def label?
         @relation_name.to_s.include?('label')
       end
+
+      def milestone?
+        @relation_name.to_s.include?('milestone')
+      end
     end
   end
 end