Add latest changes from gitlab-org/gitlab@master

5 files changed, 97 insertions, 14 deletions

diff --git a/lib/gitlab/import_export/group/tree_restorer.rb b/lib/gitlab/import_export/group/tree_restorer.rb
index f6ebd83bfaa..323e6727a9f 100644
--- a/lib/gitlab/import_export/group/tree_restorer.rb
+++ b/lib/gitlab/import_export/group/tree_restorer.rb
@@ -20,6 +20,7 @@ module Gitlab
         def restore
           @group_attributes = relation_reader.consume_attributes(nil)
           @group_members = relation_reader.consume_relation(nil, 'members')
+            .map(&:first)
 
           # We need to remove `name` and `path` as we did consume it in previous pass
           @group_attributes.delete('name')
diff --git a/lib/gitlab/import_export/json/legacy_reader.rb b/lib/gitlab/import_export/json/legacy_reader.rb
index 57579fe9def..12d6458aedc 100644
--- a/lib/gitlab/import_export/json/legacy_reader.rb
+++ b/lib/gitlab/import_export/json/legacy_reader.rb
@@ -53,6 +53,7 @@ module Gitlab
 
         def initialize(relation_names:, allowed_path:)
           @relation_names = relation_names.map(&:to_s)
+          @consumed_relations = Set.new
 
           # This is legacy reader, to be used in transition
           # period before `.ndjson`,
@@ -81,17 +82,19 @@ module Gitlab
             raise ArgumentError, "Invalid #{importable_name} passed to `consume_relation`. Use #{@allowed_path} instead."
           end
 
-          value = relations.delete(key)
+          Enumerator.new do |documents|
+            next unless @consumed_relations.add?("#{importable_path}/#{key}")
 
-          return value unless block_given?
-          return if value.nil?
+            value = relations.delete(key)
+            next if value.nil?
 
-          if value.is_a?(Array)
-            value.each.with_index do |item, idx|
-              yield(item, idx)
+            if value.is_a?(Array)
+              value.each.with_index do |item, idx|
+                documents << [item, idx]
+              end
+            else
+              documents << [value, 0]
             end
-          else
-            yield(value, 0)
           end
         end
 
diff --git a/lib/gitlab/import_export/json/ndjson_reader.rb b/lib/gitlab/import_export/json/ndjson_reader.rb
new file mode 100644
index 00000000000..e9b05afc7d4
--- /dev/null
+++ b/lib/gitlab/import_export/json/ndjson_reader.rb
@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module ImportExport
+    module JSON
+      class NdjsonReader
+        MAX_JSON_DOCUMENT_SIZE = 50.megabytes
+
+        attr_reader :dir_path
+
+        def initialize(dir_path)
+          @dir_path = dir_path
+          @consumed_relations = Set.new
+        end
+
+        def exist?
+          Dir.exist?(@dir_path)
+        end
+
+        # This can be removed once legacy_reader is deprecated.
+        def legacy?
+          false
+        end
+
+        def consume_attributes(importable_path)
+          # This reads from `tree/project.json`
+          path = file_path("#{importable_path}.json")
+          data = File.read(path, MAX_JSON_DOCUMENT_SIZE)
+          json_decode(data)
+        end
+
+        def consume_relation(importable_path, key)
+          Enumerator.new do |documents|
+            next unless @consumed_relations.add?("#{importable_path}/#{key}")
+
+            # This reads from `tree/project/merge_requests.ndjson`
+            path = file_path(importable_path, "#{key}.ndjson")
+            next unless File.exist?(path)
+
+            File.foreach(path, MAX_JSON_DOCUMENT_SIZE).with_index do |line, line_num|
+              documents << [json_decode(line), line_num]
+            end
+          end
+        end
+
+        private
+
+        def json_decode(string)
+          ActiveSupport::JSON.decode(string)
+        rescue ActiveSupport::JSON.parse_error => e
+          Gitlab::ErrorTracking.log_exception(e)
+          raise Gitlab::ImportExport::Error, 'Incorrect JSON format'
+        end
+
+        def file_path(*path)
+          File.join(dir_path, *path)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/import_export/project/tree_restorer.rb b/lib/gitlab/import_export/project/tree_restorer.rb
index 99e57d9decd..ad3720b56be 100644
--- a/lib/gitlab/import_export/project/tree_restorer.rb
+++ b/lib/gitlab/import_export/project/tree_restorer.rb
@@ -17,8 +17,13 @@ module Gitlab
         end
 
         def restore
+          unless relation_reader
+            raise Gitlab::ImportExport::Error, 'invalid import format'
+          end
+
           @project_attributes = relation_reader.consume_attributes(importable_path)
           @project_members = relation_reader.consume_relation(importable_path, 'project_members')
+            .map(&:first)
 
           if relation_tree_restorer.restore
             import_failure_service.with_retry(action: 'set_latest_merge_request_diff_ids!') do
@@ -38,14 +43,27 @@ module Gitlab
 
         def relation_reader
           strong_memoize(:relation_reader) do
-            ImportExport::JSON::LegacyReader::File.new(
-              File.join(shared.export_path, 'project.json'),
-              relation_names: reader.project_relation_names,
-              allowed_path: importable_path
-            )
+            [ndjson_relation_reader, legacy_relation_reader]
+              .compact.find(&:exist?)
           end
         end
 
+        def ndjson_relation_reader
+          return unless Feature.enabled?(:project_import_ndjson, project.namespace)
+
+          ImportExport::JSON::NdjsonReader.new(
+            File.join(shared.export_path, 'tree')
+          )
+        end
+
+        def legacy_relation_reader
+          ImportExport::JSON::LegacyReader::File.new(
+            File.join(shared.export_path, 'project.json'),
+            relation_names: reader.project_relation_names,
+            allowed_path: importable_path
+          )
+        end
+
         def relation_tree_restorer
           @relation_tree_restorer ||= RelationTreeRestorer.new(
             user: @user,
diff --git a/lib/gitlab/import_export/relation_tree_restorer.rb b/lib/gitlab/import_export/relation_tree_restorer.rb
index 78ed365cea0..056945d0294 100644
--- a/lib/gitlab/import_export/relation_tree_restorer.rb
+++ b/lib/gitlab/import_export/relation_tree_restorer.rb
@@ -67,7 +67,7 @@ module Gitlab
       end
 
       def process_relation!(relation_key, relation_definition)
-        @relation_reader.consume_relation(@importable_path, relation_key) do |data_hash, relation_index|
+        @relation_reader.consume_relation(@importable_path, relation_key).each do |data_hash, relation_index|
           process_relation_item!(relation_key, relation_definition, relation_index, data_hash)
         end
       end