1 files changed, 33 insertions, 0 deletions
diff --git a/lib/gitlab/import_export/project/sample/relation_tree_restorer.rb b/lib/gitlab/import_export/project/sample/relation_tree_restorer.rb
new file mode 100644
index 00000000000..44ccb67a531
--- /dev/null
+++ b/lib/gitlab/import_export/project/sample/relation_tree_restorer.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module ImportExport
+    module Project
+      module Sample
+        class RelationTreeRestorer < ImportExport::RelationTreeRestorer
+          def initialize(*args)
+            super
+
+            @date_calculator = Gitlab::ImportExport::Project::Sample::DateCalculator.new(dates)
+          end
+
+          private
+
+          def relation_factory_params(*args)
+            super.merge(date_calculator: @date_calculator)
+          end
+
+          def dates
+            return [] if relation_reader.legacy?
+
+            RelationFactory::DATE_MODELS.flat_map do |tag|
+              relation_reader.consume_relation(@importable_path, tag, mark_as_consumed: false).map do |model|
+                model.first['due_date']
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end