1 files changed, 13 insertions, 9 deletions
diff --git a/app/assets/javascripts/pipelines/components/parsing_utils.js b/app/assets/javascripts/pipelines/components/parsing_utils.js
index fa7330ce890..cae4e11c13f 100644
--- a/app/assets/javascripts/pipelines/components/parsing_utils.js
+++ b/app/assets/javascripts/pipelines/components/parsing_utils.js
@@ -1,5 +1,6 @@
 import { memoize } from 'lodash';
 import { createNodeDict } from '../utils';
+import { EXPLICIT_NEEDS_PROPERTY, NEEDS_PROPERTY } from '../constants';
 import { createSankey } from './dag/drawing_utils';
 
 /*
@@ -15,12 +16,14 @@ const deduplicate = (item, itemIndex, arr) => {
   return foundIdx === itemIndex;
 };
 
-export const makeLinksFromNodes = (nodes, nodeDict) => {
+export const makeLinksFromNodes = (nodes, nodeDict, { needsKey = NEEDS_PROPERTY } = {}) => {
   const constantLinkValue = 10; // all links are the same weight
   return nodes
     .map(({ jobs, name: groupName }) =>
-      jobs.map(({ needs = [] }) =>
-        needs.reduce((acc, needed) => {
+      jobs.map((job) => {
+        const needs = job[needsKey] || [];
+
+        return needs.reduce((acc, needed) => {
           // It's possible that we have an optional job, which
           // is being needed by another job. In that scenario,
           // the needed job doesn't exist, so we don't want to
@@ -34,8 +37,8 @@ export const makeLinksFromNodes = (nodes, nodeDict) => {
           }
 
           return acc;
-        }, []),
-      ),
+        }, []);
+      }),
     )
     .flat(2);
 };
@@ -76,9 +79,9 @@ export const filterByAncestors = (links, nodeDict) =>
     return !allAncestors.includes(source);
   });
 
-export const parseData = (nodes) => {
-  const nodeDict = createNodeDict(nodes);
-  const allLinks = makeLinksFromNodes(nodes, nodeDict);
+export const parseData = (nodes, { needsKey = NEEDS_PROPERTY } = {}) => {
+  const nodeDict = createNodeDict(nodes, { needsKey });
+  const allLinks = makeLinksFromNodes(nodes, nodeDict, { needsKey });
   const filteredLinks = allLinks.filter(deduplicate);
   const links = filterByAncestors(filteredLinks, nodeDict);
 
@@ -123,7 +126,8 @@ export const removeOrphanNodes = (sankeyfiedNodes) => {
 export const listByLayers = ({ stages }) => {
   const arrayOfJobs = stages.flatMap(({ groups }) => groups);
   const parsedData = parseData(arrayOfJobs);
-  const dataWithLayers = createSankey()(parsedData);
+  const explicitParsedData = parseData(arrayOfJobs, { needsKey: EXPLICIT_NEEDS_PROPERTY });
+  const dataWithLayers = createSankey()(explicitParsedData);
 
   const pipelineLayers = dataWithLayers.nodes.reduce((acc, { layer, name }) => {
     /* sort groups by layer */