1 files changed, 71 insertions, 17 deletions

diff --git a/lib/gitlab/import_export/json/streaming_serializer.rb b/lib/gitlab/import_export/json/streaming_serializer.rb
index 05b7679e0ff..ec42c5e51c0 100644
--- a/lib/gitlab/import_export/json/streaming_serializer.rb
+++ b/lib/gitlab/import_export/json/streaming_serializer.rb
@@ -38,16 +38,6 @@ module Gitlab
           end
         end
 
-        private
-
-        attr_reader :json_writer, :relations_schema, :exportable
-
-        def serialize_root
-          attributes = exportable.as_json(
-            relations_schema.merge(include: nil, preloads: nil))
-          json_writer.write_attributes(@exportable_path, attributes)
-        end
-
         def serialize_relation(definition)
           raise ArgumentError, 'definition needs to be Hash' unless definition.is_a?(Hash)
           raise ArgumentError, 'definition needs to have exactly one Hash element' unless definition.one?
@@ -64,17 +54,22 @@ module Gitlab
           end
         end
 
+        private
+
+        attr_reader :json_writer, :relations_schema, :exportable
+
+        def serialize_root
+          attributes = exportable.as_json(
+            relations_schema.merge(include: nil, preloads: nil))
+          json_writer.write_attributes(@exportable_path, attributes)
+        end
+
         def serialize_many_relations(key, records, options)
           enumerator = Enumerator.new do |items|
             key_preloads = preloads&.dig(key)
-            records = records.preload(key_preloads) if key_preloads
 
-            records.in_batches(of: batch_size) do |batch| # rubocop:disable Cop/InBatches
-              # order each batch by its primary key to ensure
-              # consistent and predictable ordering of each exported relation
-              # as additional `WHERE` clauses can impact the order in which data is being
-              # returned by database when no `ORDER` is specified
-              batch = batch.reorder(batch.klass.primary_key)
+            batch(records, key) do |batch|
+              batch = batch.preload(key_preloads) if key_preloads
 
               batch.each do |record|
                 items << Raw.new(record.to_json(options))
@@ -85,6 +80,29 @@ module Gitlab
           json_writer.write_relation_array(@exportable_path, key, enumerator)
         end
 
+        def batch(relation, key)
+          opts = { of: batch_size }
+          order_by = reorders(relation, key)
+
+          # we need to sort issues by non primary key column(relative_position)
+          # and `in_batches` does not support that
+          if order_by
+            scope = relation.reorder(order_by)
+
+            Gitlab::Pagination::Keyset::Iterator.new(scope: scope, use_union_optimization: true).each_batch(**opts) do |batch|
+              yield batch
+            end
+          else
+            relation.in_batches(**opts) do |batch| # rubocop:disable Cop/InBatches
+              # order each batch by its primary key to ensure
+              # consistent and predictable ordering of each exported relation
+              # as additional `WHERE` clauses can impact the order in which data is being
+              # returned by database when no `ORDER` is specified
+              yield batch.reorder(batch.klass.primary_key)
+            end
+          end
+        end
+
         def serialize_many_each(key, records, options)
           enumerator = Enumerator.new do |items|
             records.each do |record|
@@ -112,6 +130,42 @@ module Gitlab
         def batch_size
           @batch_size ||= self.class.batch_size(@exportable)
         end
+
+        def reorders(relation, key)
+          export_reorder = relations_schema[:export_reorder]&.dig(key)
+          return unless export_reorder
+
+          custom_reorder(relation.klass, export_reorder)
+        end
+
+        def custom_reorder(klass, order_by)
+          arel_table = klass.arel_table
+          column = order_by[:column] || klass.primary_key
+          direction = order_by[:direction] || :asc
+          nulls_position = order_by[:nulls_position] || :nulls_last
+
+          arel_order_classes = ::Gitlab::Pagination::Keyset::ColumnOrderDefinition::AREL_ORDER_CLASSES.invert
+          reverse_direction = ::Gitlab::Pagination::Keyset::ColumnOrderDefinition::REVERSED_ORDER_DIRECTIONS[direction]
+          reverse_nulls_position = ::Gitlab::Pagination::Keyset::ColumnOrderDefinition::REVERSED_NULL_POSITIONS[nulls_position]
+          order_expression = ::Gitlab::Database.nulls_order(column, direction, nulls_position)
+          reverse_order_expression = ::Gitlab::Database.nulls_order(column, reverse_direction, reverse_nulls_position)
+
+          ::Gitlab::Pagination::Keyset::Order.build([
+            ::Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+              attribute_name: column,
+              column_expression: arel_table[column],
+              order_expression: order_expression,
+              reversed_order_expression: reverse_order_expression,
+              order_direction: direction,
+              nullable: nulls_position,
+              distinct: false
+            ),
+            ::Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+              attribute_name: klass.primary_key,
+              order_expression: arel_order_classes[direction].new(arel_table[klass.primary_key.to_sym])
+            )
+          ])
+        end
       end
     end
   end