From 0500e50463d27dba43b188f7179be6b50ebdb82d Mon Sep 17 00:00:00 2001
From: Gabe Lyons <itsgabelyons@gmail.com>
Date: Mon, 5 Apr 2021 12:15:05 -0700
Subject: [PATCH] feat(lineage): adding directionality to lineage edges to make
 the visualization more clear (#2335)

---
 .../src/app/lineage/LineageEntityNode.tsx        |  4 ++--
 .../lineage/LineageTreeNodeAndEdgeRenderer.tsx   |  2 ++
 datahub-web-react/src/app/lineage/LineageViz.tsx | 16 ++++++++++++++--
 .../src/app/lineage/utils/adjustVXTreeLayout.ts  |  6 ++++--
 4 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/datahub-web-react/src/app/lineage/LineageEntityNode.tsx b/datahub-web-react/src/app/lineage/LineageEntityNode.tsx
index 0a404c6504..5e9150b5c2 100644
--- a/datahub-web-react/src/app/lineage/LineageEntityNode.tsx
+++ b/datahub-web-react/src/app/lineage/LineageEntityNode.tsx
@@ -13,8 +13,8 @@ function truncate(input, length) {
     return input;
 }
 
-const width = 140;
-const height = 80;
+export const width = 140;
+export const height = 80;
 const centerX = -width / 2;
 const centerY = -height / 2;
 
diff --git a/datahub-web-react/src/app/lineage/LineageTreeNodeAndEdgeRenderer.tsx b/datahub-web-react/src/app/lineage/LineageTreeNodeAndEdgeRenderer.tsx
index 93c2cf6a35..350d164e85 100644
--- a/datahub-web-react/src/app/lineage/LineageTreeNodeAndEdgeRenderer.tsx
+++ b/datahub-web-react/src/app/lineage/LineageTreeNodeAndEdgeRenderer.tsx
@@ -119,6 +119,8 @@ export default function LineageTreeNodeAndEdgeRenderer({
                         fill="none"
                         key={`edge-${link.source.data.urn}-${link.target.data.urn}-${direction}`}
                         data-testid={`edge-${link.source.data.urn}-${link.target.data.urn}-${direction}`}
+                        markerEnd="url(#triangle-downstream)"
+                        markerStart="url(#triangle-upstream)"
                     />
                 );
             })}
diff --git a/datahub-web-react/src/app/lineage/LineageViz.tsx b/datahub-web-react/src/app/lineage/LineageViz.tsx
index c50341ceb9..dc2a7c6e6d 100644
--- a/datahub-web-react/src/app/lineage/LineageViz.tsx
+++ b/datahub-web-react/src/app/lineage/LineageViz.tsx
@@ -92,9 +92,9 @@ export default function LineageViz({
                     >
                         <defs>
                             <marker
-                                id="triangle"
+                                id="triangle-downstream"
                                 viewBox="0 0 10 10"
-                                refX="80"
+                                refX="10"
                                 refY="5"
                                 markerUnits="strokeWidth"
                                 markerWidth="10"
@@ -103,6 +103,18 @@ export default function LineageViz({
                             >
                                 <path d="M 0 0 L 10 5 L 0 10 z" fill="#000" />
                             </marker>
+                            <marker
+                                id="triangle-upstream"
+                                viewBox="0 0 10 10"
+                                refX="0"
+                                refY="5"
+                                markerUnits="strokeWidth"
+                                markerWidth="10"
+                                markerHeight="10"
+                                orient="auto"
+                            >
+                                <path d="M 0 5 L 10 10 L 10 0 L 0 5 z" fill="#000" />
+                            </marker>
                         </defs>
                         <rect width={width} height={height} fill="#f6f8fa" />
                         <LineageTree
diff --git a/datahub-web-react/src/app/lineage/utils/adjustVXTreeLayout.ts b/datahub-web-react/src/app/lineage/utils/adjustVXTreeLayout.ts
index 57e8742882..22c7f088ee 100644
--- a/datahub-web-react/src/app/lineage/utils/adjustVXTreeLayout.ts
+++ b/datahub-web-react/src/app/lineage/utils/adjustVXTreeLayout.ts
@@ -1,5 +1,7 @@
 import { HierarchyPointNode } from '@vx/hierarchy/lib/types';
 import { NodeData, Direction } from '../types';
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+import { width as nodeWidth } from '../LineageEntityNode';
 
 export default function adjustVXTreeLayout({
     tree,
@@ -50,12 +52,12 @@ export default function adjustVXTreeLayout({
     const edgesToReturn = tree.links().map((linkToCopy) => ({
         target: {
             x: linkToCopy.target.x,
-            y: linkToCopy.target.y,
+            y: linkToCopy.target.y + (direction === Direction.Upstream ? 0 : -(nodeWidth / 2)),
             data: { ...linkToCopy.target.data },
         },
         source: {
             x: linkToCopy.source.x,
-            y: linkToCopy.source.y,
+            y: linkToCopy.source.y + (direction === Direction.Upstream ? -(nodeWidth / 2) : 0),
             data: { ...linkToCopy.source.data },
         },
     }));