Improve visual display for lineage (#2753)

* Feedback updates for lineage graph * Making display better. * Removing log. --------- Co-authored-by: phix <peter.hicks@astronomer.io>
MarquezProject · Feb 22, 2024 · c756cb2 · c756cb2
1 parent ce2d12a
commit c756cb2
Show file tree

Hide file tree

Showing 6 changed files with 177 additions and 94 deletions.
diff --git a/web/src/components/core/tooltip/MQTooltip.tsx b/web/src/components/core/tooltip/MQTooltip.tsx
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { createTheme } from '@mui/material/styles'
+import { grey } from '@mui/material/colors'
 import { useTheme } from '@emotion/react'
 import React, { ReactElement } from 'react'
 import Tooltip from '@mui/material/Tooltip'
@@ -21,8 +22,8 @@ const MQTooltip: React.FC<MqToolTipProps> = ({ title, children, placement }) =>
       componentsProps={{
         tooltip: {
           sx: {
-            backgroundColor: theme.palette.background.default,
-            color: theme.palette.common.white,
+            backgroundColor: `${theme.palette.common.white}`,
+            color: grey['900'],
             border: `1px solid ${theme.palette.common.white}`,
             maxWidth: '600px',
             fontSize: 14,

diff --git a/web/src/helpers/theme.ts b/web/src/helpers/theme.ts
@@ -43,7 +43,7 @@ export const theme = createTheme({
       main: '#7D7D7D',
     },
     info: {
-      main: '#FECC00',
+      main: '#9c98ec',
     },
     background: {
       default: '#191f26',

diff --git a/web/src/routes/table-level/TableLevel.tsx b/web/src/routes/table-level/TableLevel.tsx
@@ -46,6 +46,8 @@ const ColumnLevel: React.FC<ColumnLevelProps> = ({
 
   const graphControls = useRef<ZoomPanControls>()
 
+  const collapsedNodes = searchParams.get('collapsedNodes')
+
   useEffect(() => {
     if (name && namespace && nodeType) {
       fetchLineage(nodeType as JobOrDataset, namespace, name, depth)
@@ -79,7 +81,8 @@ const ColumnLevel: React.FC<ColumnLevelProps> = ({
     lineage,
     `${nodeType}:${namespace}:${name}`,
     isCompact,
-    isFull
+    isFull,
+    collapsedNodes
   )
 
   useEffect(() => {

diff --git a/web/src/routes/table-level/TableLineageDatasetNode.tsx b/web/src/routes/table-level/TableLineageDatasetNode.tsx
@@ -1,3 +1,4 @@
+import { ChevronLeft } from '@mui/icons-material'
 import { FontAwesomeIcon } from '@fortawesome/react-fontawesome'
 import { IState } from '../../store/reducers'
 import { LineageDataset } from '../../components/lineage/types'
@@ -6,11 +7,12 @@ import { PositionedNode } from '../../../libs/graph'
 import { THEME_EXTRA, theme } from '../../helpers/theme'
 import { TableLineageDatasetNodeData } from './nodes'
 import { connect } from 'react-redux'
+
 import { faDatabase } from '@fortawesome/free-solid-svg-icons/faDatabase'
-import { grey } from '@mui/material/colors'
 import { truncateText } from '../../helpers/text'
-import { useNavigate, useParams } from 'react-router-dom'
+import { useNavigate, useParams, useSearchParams } from 'react-router-dom'
 import Box from '@mui/system/Box'
+import IconButton from '@mui/material/IconButton'
 import MQTooltip from '../../components/core/tooltip/MQTooltip'
 import React from 'react'
 
@@ -31,6 +33,8 @@ const TableLineageDatasetNode = ({ node }: TableLineageDatasetNodeProps & StateP
   const navigate = useNavigate()
   const { name, namespace } = useParams()
   const isSelected = name === node.data.dataset.name && namespace === node.data.dataset.namespace
+  const [searchParams, setSearchParams] = useSearchParams()
+  const isCollapsed = searchParams.get('collapsedNodes')?.split(',').includes(node.id)
 
   const handleClick = () => {
     navigate(
@@ -57,56 +61,109 @@ const TableLineageDatasetNode = ({ node }: TableLineageDatasetNodeProps & StateP
   }
 
   return (
-    <MQTooltip title={addToToolTip(node.data.dataset)}>
-      <g>
-        <Box
-          component={'rect'}
-          sx={{
-            x: 0,
-            y: 0,
-            width: node.width,
-            height: node.height,
-            stroke: isSelected ? theme.palette.primary.main : grey['100'],
-            filter: isSelected ? `drop-shadow( 0 0 4px ${theme.palette.primary.main})` : 'none',
-            rx: 4,
-            fill: theme.palette.background.paper,
-            cursor: 'pointer',
-            transition: 'filter 0.3',
-          }}
-          cursor={'pointer'}
-          onClick={handleClick}
-        />
-        <FontAwesomeIcon
-          aria-hidden={'true'}
-          title={'Job'}
-          icon={faDatabase}
-          width={ICON_SIZE}
-          height={ICON_SIZE}
-          x={4}
-          y={ICON_SIZE / 2}
-          color={theme.palette.warning.main}
-          cursor={'pointer'}
-          onClick={handleClick}
-        />
-        <text fontSize='8' fill={'white'} x={20} y={14} cursor={'pointer'} onClick={handleClick}>
-          {truncateText(node.data.dataset.name, 15)}
-        </text>
-        {!isCompact &&
-          node.data.dataset.fields.map((field, index) => {
-            return (
-              <text
-                key={field.name}
-                fontSize='8'
-                fill={THEME_EXTRA.typography.subdued}
-                x={20}
-                y={14 + 10 + 10 * (index + 1)}
-              >
-                - {truncateText(field.name, 15)}
-              </text>
-            )
-          })}
-      </g>
-    </MQTooltip>
+    <g>
+      <Box
+        component={'rect'}
+        sx={{
+          x: 0,
+          y: 0,
+          width: node.width,
+          height: node.height,
+          filter: isSelected ? `drop-shadow( 0 0 4px ${theme.palette.primary.main})` : 'none',
+          rx: 4,
+          fill: theme.palette.background.paper,
+          cursor: 'pointer',
+          transition: 'all 0.3',
+        }}
+        cursor={'pointer'}
+        onClick={handleClick}
+      />
+      <Box
+        component={'rect'}
+        x={0}
+        y={0}
+        height={24}
+        width={24}
+        sx={{ rx: 4, fill: theme.palette.info.main }}
+      />
+
+      <FontAwesomeIcon
+        aria-hidden={'true'}
+        title={'Job'}
+        icon={faDatabase}
+        width={ICON_SIZE}
+        height={ICON_SIZE}
+        x={6}
+        y={ICON_SIZE / 2}
+        color={theme.palette.common.white}
+        onClick={handleClick}
+      />
+      <foreignObject width={16} height={24} x={node.width - 18} y={0}>
+        <MQTooltip title={isCollapsed ? 'Expand' : 'Collapse'} placement={'top'}>
+          <IconButton
+            sx={{ width: 10, height: 10 }}
+            onClick={(event) => {
+              event.stopPropagation()
+              const collapsedNodes = searchParams.get('collapsedNodes')
+              if (collapsedNodes) {
+                const collapsedNodesArray = collapsedNodes.split(',')
+                if (collapsedNodesArray.includes(node.id)) {
+                  collapsedNodesArray.splice(collapsedNodesArray.indexOf(node.id), 1)
+                } else {
+                  collapsedNodesArray.push(node.id)
+                }
+                searchParams.set('collapsedNodes', collapsedNodesArray.toString())
+              } else {
+                searchParams.set('collapsedNodes', node.id)
+              }
+              setSearchParams(searchParams)
+            }}
+          >
+            <ChevronLeft
+              sx={{
+                width: 10,
+                height: 10,
+                rotate: !isCollapsed ? '-90deg' : 0,
+                transition: 'rotate .3s',
+              }}
+            />
+          </IconButton>
+        </MQTooltip>
+      </foreignObject>
+      <MQTooltip title={addToToolTip(node.data.dataset)}>
+        <g>
+          <text
+            fontSize='8'
+            fontFamily={`${'Source Code Pro'}, mono`}
+            fill={'white'}
+            x={28}
+            y={10}
+            onClick={handleClick}
+            cursor={'pointer'}
+          >
+            DATASET
+          </text>
+          <text fontSize='8' fill={'white'} x={28} y={20} cursor={'pointer'} onClick={handleClick}>
+            {truncateText(node.data.dataset.name, 15)}
+          </text>
+        </g>
+      </MQTooltip>
+
+      {!isCompact &&
+        node.data.dataset.fields.map((field, index) => {
+          return (
+            <text
+              key={field.name}
+              fontSize='8'
+              fill={THEME_EXTRA.typography.subdued}
+              x={10}
+              y={14 + 10 + 10 * (index + 1)}
+            >
+              - {truncateText(field.name, 20)}
+            </text>
+          )
+        })}
+    </g>
   )
 }
 

diff --git a/web/src/routes/table-level/TableLineageJobNode.tsx b/web/src/routes/table-level/TableLineageJobNode.tsx
@@ -6,7 +6,6 @@ import { PositionedNode } from '../../../libs/graph'
 import { TableLineageJobNodeData } from './nodes'
 import { connect } from 'react-redux'
 import { faCog } from '@fortawesome/free-solid-svg-icons/faCog'
-import { grey } from '@mui/material/colors'
 import { theme } from '../../helpers/theme'
 import { truncateText } from '../../helpers/text'
 import { useNavigate, useParams } from 'react-router-dom'
@@ -53,41 +52,60 @@ const TableLineageJobNode = ({ node }: TableLineageJobNodeProps & StateProps) =>
   }
 
   return (
-    <MQTooltip title={addToToolTip(node.data.job)}>
-      <g>
-        <Box
-          component={'rect'}
-          sx={{
-            x: 0,
-            y: 0,
-            width: node.width,
-            height: node.height,
-            stroke: isSelected ? theme.palette.primary.main : grey['100'],
-            filter: isSelected ? `drop-shadow( 0 0 4px ${theme.palette.primary.main})` : 'none',
-            rx: 4,
-            fill: theme.palette.background.paper,
-            cursor: 'pointer',
-            transition: 'filter 0.3',
-          }}
-          onClick={handleClick}
-        />
-        <FontAwesomeIcon
-          aria-hidden={'true'}
-          title={'Job'}
-          icon={faCog}
-          width={ICON_SIZE}
-          height={ICON_SIZE}
-          x={4}
-          y={ICON_SIZE / 2}
-          color={theme.palette.primary.main}
-          cursor={'pointer'}
-          onClick={handleClick}
-        />
-        <text fontSize='8' fill={'white'} x={20} y={14} onClick={handleClick} cursor={'pointer'}>
-          {truncateText(node.data.job.name, 15)}
-        </text>
-      </g>
-    </MQTooltip>
+    <g>
+      <Box
+        component={'rect'}
+        sx={{
+          x: 0,
+          y: 0,
+          width: node.width,
+          height: node.height,
+          filter: isSelected ? `drop-shadow( 0 0 4px ${theme.palette.primary.main})` : 'none',
+          rx: 4,
+          fill: theme.palette.background.paper,
+          cursor: 'pointer',
+          transition: 'filter 0.3',
+        }}
+        onClick={handleClick}
+      />
+      <Box
+        component={'rect'}
+        x={0}
+        y={0}
+        height={node.height}
+        width={24}
+        sx={{ rx: 4, fill: theme.palette.primary.main }}
+      />
+      <FontAwesomeIcon
+        aria-hidden={'true'}
+        title={'Job'}
+        icon={faCog}
+        width={ICON_SIZE}
+        height={ICON_SIZE}
+        x={6}
+        y={ICON_SIZE / 2}
+        color={theme.palette.common.white}
+        onClick={handleClick}
+      />
+      <MQTooltip title={addToToolTip(node.data.job)}>
+        <g>
+          <text
+            fontSize='8'
+            fontFamily={`${'Source Code Pro'}, mono`}
+            fill={'white'}
+            x={28}
+            y={10}
+            onClick={handleClick}
+            cursor={'pointer'}
+          >
+            JOB
+          </text>
+          <text fontSize='8' fill={'white'} x={28} y={20} onClick={handleClick} cursor={'pointer'}>
+            {truncateText(node.data.job.name, 16)}
+          </text>
+        </g>
+      </MQTooltip>
+    </g>
   )
 }
 

diff --git a/web/src/routes/table-level/layout.ts b/web/src/routes/table-level/layout.ts
@@ -76,14 +76,17 @@ export const createElkNodes = (
   lineageGraph: LineageGraph,
   currentGraphNode: Nullable<string>,
   isCompact: boolean,
-  isFull: boolean
+  isFull: boolean,
+  collapsedNodes: Nullable<string>
 ) => {
   const downstreamNodes = findDownstreamNodes(lineageGraph, currentGraphNode)
   const upstreamNodes = findUpstreamNodes(lineageGraph, currentGraphNode)
 
   const nodes: ElkNode<JobOrDataset, TableLevelNodeData>[] = []
   const edges: Edge[] = []
 
+  const collapsedNodesAsArray = collapsedNodes?.split(',')
+
   const filteredGraph = lineageGraph.graph.filter((node) => {
     if (isFull) return true
     return (
@@ -124,7 +127,8 @@ export const createElkNodes = (
         id: node.id,
         kind: node.type,
         width: 112,
-        height: isCompact ? 24 : 34 + data.fields.length * 10,
+        height:
+          isCompact || collapsedNodesAsArray?.includes(node.id) ? 24 : 34 + data.fields.length * 10,
         data: {
           dataset: data,
         },