1 files changed, 31 insertions, 0 deletions
diff --git a/lib/gitlab/ci/pipeline/chain/assign_partition.rb b/lib/gitlab/ci/pipeline/chain/assign_partition.rb
new file mode 100644
index 00000000000..4b8efe13d44
--- /dev/null
+++ b/lib/gitlab/ci/pipeline/chain/assign_partition.rb
@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Ci
+    module Pipeline
+      module Chain
+        class AssignPartition < Chain::Base
+          include Chain::Helpers
+
+          def perform!
+            @pipeline.partition_id = find_partition_id
+          end
+
+          def break?
+            @pipeline.errors.any?
+          end
+
+          private
+
+          def find_partition_id
+            if @command.creates_child_pipeline?
+              @command.parent_pipeline_partition_id
+            else
+              ::Ci::Pipeline.current_partition_value
+            end
+          end
+        end
+      end
+    end
+  end
+end