refactor code based on feedback

lib/gitlab/import_export/group_project_object_builder.rb

@@ -36,19 +36,17 @@ module Gitlab
       end
 
       def where_clause
-        return { project_id: @project.id } unless milestone? || label?
-
         @attributes.slice('title').map do |key, value|
           if @group
-            project_clause(key, value).or(group_clause(key, value))
+            project_group_clause(key, value)
           else
             project_clause(key, value)
           end
         end.reduce(:or)
       end
 
-      def group_clause(key, value)
+      def project_group_clause(key, value)
-        table[key].eq(value).and(table[:group_id].eq(@group.id))
+        table[key].eq(value).and(table[:project_id].eq(@project.id).or(table[:group_id].eq(@group.id)))
       end
 
       def project_clause(key, value)
@@ -92,11 +90,17 @@ module Gitlab
         # we set the IID as the maximum. The rest of them are fixed.
         group_milestone = @project.milestones.find_by(iid: @attributes['iid'])
 
-        group_milestone.update!(iid: max_milestone_iid + 1) if group_milestone
+        group_milestone.update!(iid:  max_milestone_iid(group_milestone)) if group_milestone
       end
 
-      def max_milestone_iid
+      def max_milestone_iid(group_milestone)
-        [@attributes['iid'], @project.milestones.maximum(:iid)].compact.max
+        init_iid = [@attributes['iid'], @project.milestones.maximum(:iid)].compact.max + 1
+
+        InternalId::InternalIdGenerator.new(group_milestone,
+                                            { project: @project },
+                                            :milestones,
+                                            init_iid
+        ).generate
       end
     end
   end

lib/gitlab/import_export/project_tree_restorer.rb

@@ -1,7 +1,7 @@
 module Gitlab
   module ImportExport
     class ProjectTreeRestorer
-      # Relations which cannot be saved at project level
+      # Relations which cannot be saved at project level (and have a group assigned)
       GROUP_MODELS = [GroupLabel, Milestone].freeze
 
       def initialize(user:, shared:, project:)
@@ -83,7 +83,7 @@ module Gitlab
       # For example, in the case of an existing group label that matched the title.
       def remove_group_models(relation_hash)
         relation_hash.reject! do |value|
-          value.respond_to?(:group_id) && value.group_id && GROUP_MODELS.include?(value.class)
+          GROUP_MODELS.include?(value.class) && value.group_id
         end
       end
 

lib/gitlab/import_export/relation_factory.rb

@@ -150,16 +150,16 @@ module Gitlab
       end
 
       def update_project_references
-        project_id = @relation_hash.delete('project_id')
-
         # If source and target are the same, populate them with the new project ID.
         if @relation_hash['source_project_id']
-          @relation_hash['source_project_id'] = same_source_and_target? ? project_id : MergeRequestParser::FORKED_PROJECT_ID
+          @relation_hash['source_project_id'] = same_source_and_target? ? @relation_hash['project_id'] : MergeRequestParser::FORKED_PROJECT_ID
         end
 
-        # project_id may not be part of the export, but we always need to populate it if required.
+        @relation_hash['target_project_id'] = @relation_hash['project_id'] if @relation_hash['target_project_id']
-        @relation_hash['project_id'] = project_id
+      end
-        @relation_hash['target_project_id'] = project_id if @relation_hash['target_project_id']
+
+      def same_source_and_target?
+        @relation_hash['target_project_id'] && @relation_hash['target_project_id'] == @relation_hash['source_project_id']
       end
 
       def update_group_references
@@ -169,10 +169,6 @@ module Gitlab
         @relation_hash['group_id'] = @project.group&.id
       end
 
-      def same_source_and_target?
-        @relation_hash['target_project_id'] && @relation_hash['target_project_id'] == @relation_hash['source_project_id']
-      end
-
       def reset_tokens!
         return unless Gitlab::ImportExport.reset_tokens? && TOKEN_RESET_MODELS.include?(@relation_name.to_s)
 
@@ -267,6 +263,8 @@ module Gitlab
       end
 
       def find_or_create_object!
+        return relation_class.find_or_create_by(project_id: @project.id) if @relation_name == :project_feature
+
         # Can't use IDs as validation exists calling `group` or `project` attributes
         finder_hash = parsed_relation_hash.tap do |hash|
           hash[:group] = @project.group if relation_class.attribute_method?('group_id')