refactoring relation factory, changed from module to class

2 files changed, 72 insertions, 58 deletions

diff --git a/lib/gitlab/import_export/project_tree_restorer.rb b/lib/gitlab/import_export/project_tree_restorer.rb
index 8d3a016ad71..dc1e477a82f 100644
--- a/lib/gitlab/import_export/project_tree_restorer.rb
+++ b/lib/gitlab/import_export/project_tree_restorer.rb
@@ -32,6 +32,11 @@ module Gitlab
                                                                     project: project)
       end
 
+      # Loops through the tree of models defined in import_export.yml and
+      # finds them in the imported JSON so they can be instantiated and saved
+      # in the DB. The structure and relationships between models are guessed from
+      # the configuration yaml file too.
+      # Finally, it updates each attribute in the newly imported project.
       def create_relations
         saved = []
         default_relation_list.each do |relation|
diff --git a/lib/gitlab/import_export/relation_factory.rb b/lib/gitlab/import_export/relation_factory.rb
index cecb4747822..3738dbaebd3 100644
--- a/lib/gitlab/import_export/relation_factory.rb
+++ b/lib/gitlab/import_export/relation_factory.rb
@@ -1,7 +1,6 @@
 module Gitlab
   module ImportExport
-    module RelationFactory
-      extend self
+    class RelationFactory
 
       OVERRIDES = { snippets: :project_snippets,
                     ci_commits: 'Ci::Commit',
@@ -13,47 +12,54 @@ module Gitlab
 
       USER_REFERENCES = %w[author_id assignee_id updated_by_id user_id].freeze
 
-      # Guesses a model and saves it to the DB given its name `relation_sym`
-      def create(relation_sym:, relation_hash:, members_mapper:, user_admin:)
-        relation_sym = parse_relation_sym(relation_sym)
+      def self.create(*args)
+        new(*args).create
+      end
 
-        klass = relation_class(relation_sym)
-        relation_hash.delete('id')
+      def initialize(relation_sym:, relation_hash:, members_mapper:, user_admin:)
+        @relation_name = OVERRIDES[relation_sym] || relation_sym
+        @relation_hash = relation_hash.except('id')
+        @members_mapper = members_mapper
+        @user_admin = user_admin
+      end
 
-        update_missing_author(relation_hash, members_mapper, user_admin) if relation_sym == :notes
-        update_user_references(relation_hash, members_mapper.map)
-        update_project_references(relation_hash, klass)
-        reset_tokens(relation_hash) if relation_sym == 'Ci::Trigger'
+      # Creates an object from an actual model with name "relation_sym" with params from
+      # the relation_hash, updating references with new object IDs, mapping users using
+      # the "members_mapper" object, also updating notes if required.
+      def create
+        set_note_author if @relation_name == :notes
+        update_user_references
+        update_project_references
+        reset_tokens if @relation_name == 'Ci::Trigger'
 
-        generate_imported_object(klass, relation_hash, relation_sym)
+        generate_imported_object
       end
 
       private
 
-      def update_user_references(relation_hash, members_map)
+      def update_user_references
         USER_REFERENCES.each do |reference|
-          if relation_hash[reference]
-            relation_hash[reference] = members_map[relation_hash[reference]]
+          if @relation_hash[reference]
+            @relation_hash[reference] = @members_mapper.map[@relation_hash[reference]]
           end
         end
       end
 
-      def update_missing_author(relation_hash, members_map, user_admin)
-        old_author_id = relation_hash['author_id']
+      # Sets the author for a note. If the user importing the project
+      # has admin access, an actual mapping with new project members
+      # will be used. Otherwise, a note stating the original author name
+      # is left.
+      def set_note_author
+        old_author_id = @relation_hash['author_id']
 
         # Users with admin access can map users
-        if user_admin
-          relation_hash['author_id'] = members_map.map[old_author_id]
-        else
-          relation_hash['author_id'] = members_map.default_user_id
-        end
+        @relation_hash['author_id'] = admin_user? ? @members_mapper.map[old_author_id] : @members_mapper.default_user_id
 
-        author = relation_hash.delete('author')
+        author = @relation_hash.delete('author')
 
-        return unless user_admin && members_map.note_member_list.include?(old_author_id)
-
-        relation_hash['note'] = '*Blank note*' if relation_hash['note'].blank?
-        relation_hash['note'] += missing_author_note(relation_hash['updated_at'], author['name'])
+        if admin_user? && @members_mapper.note_member_list.include?(old_author_id)
+          update_note_for_missing_author(author['name'])
+        end
       end
 
       def missing_author_note(updated_at, author_name)
@@ -61,57 +67,60 @@ module Gitlab
         "\n\n *By #{author_name} on #{timestamp} (imported from GitLab project)*"
       end
 
-      def generate_imported_object(klass, relation_hash, relation_sym)
-        if relation_sym == 'commit_status' # call #trace= method after assigning the other attributes
-          trace = relation_hash.delete('trace')
-          imported_object(klass, relation_hash) do |imported_object|
-            imported_object.trace = trace
-            imported_object.commit_id = nil
+      def generate_imported_object
+        if @relation_sym == 'commit_status' # call #trace= method after assigning the other attributes
+          trace = @relation_hash.delete('trace')
+          imported_object do |object|
+            object.trace = trace
+            object.commit_id = nil
           end
         else
-          imported_object(klass, relation_hash)
+          imported_object
         end
       end
 
-      def update_project_references(relation_hash, klass)
-        project_id = relation_hash.delete('project_id')
-
-        if relation_hash['source_project_id'] && relation_hash['target_project_id']
-          # If source and target are the same, populate them with the new project ID.
-          if relation_hash['target_project_id'] == relation_hash['source_project_id']
-            relation_hash['source_project_id'] = project_id
-          else
-            relation_hash['source_project_id'] = -1
-          end
-        end
-        relation_hash['target_project_id'] = project_id if relation_hash['target_project_id']
+      def update_project_references
+        project_id = @relation_hash.delete('project_id')
 
         # project_id may not be part of the export, but we always need to populate it if required.
-        relation_hash['project_id'] = project_id if klass.column_names.include?('project_id')
-        relation_hash['gl_project_id'] = project_id if relation_hash ['gl_project_id']
+        @relation_hash['project_id'] = project_id if relation_class.column_names.include?('project_id')
+        @relation_hash['gl_project_id'] = project_id if @relation_hash['gl_project_id']
+        @relation_hash['target_project_id'] = project_id if @relation_hash['target_project_id']
+        @relation_hash['source_project_id'] = -1 if @relation_hash['source_project_id']
+
+        # If source and target are the same, populate them with the new project ID.
+        if @relation_hash['source_project_id'] && @relation_hash['target_project_id'] &&
+          @relation_hash['target_project_id'] == @relation_hash['source_project_id']
+          @relation_hash['source_project_id'] = project_id
+        end
       end
 
-      def reset_tokens(relation_hash)
+      def reset_tokens
         return unless Gitlab::ImportExport.reset_tokens?
 
         # If we import/export a project to the same instance, tokens will have to be reseated.
-        relation_hash['token'] = nil
-      end
-
-      def relation_class(relation_sym)
-        relation_sym.to_s.classify.constantize
+        @relation_hash['token'] = nil
       end
 
-      def parse_relation_sym(relation_sym)
-        OVERRIDES[relation_sym] || relation_sym
+      def relation_class
+        @relation_class ||= @relation_name.to_s.classify.constantize
       end
 
-      def imported_object(klass, relation_hash)
-        imported_object = klass.new(relation_hash)
+      def imported_object
+        imported_object = relation_class.new(@relation_hash)
         yield(imported_object) if block_given?
         imported_object.importing = true if imported_object.respond_to?(:importing)
         imported_object
       end
+
+      def update_note_for_missing_author(author_name)
+        @relation_hash['note'] = '*Blank note*' if @relation_hash['note'].blank?
+        @relation_hash['note'] += missing_author_note(@relation_hash['updated_at'], author_name)
+      end
+
+      def admin_user?
+        @user_admin
+      end
     end
   end
 end