From 7adaeb3048f9c7555655c011ee9df3db1e385922 Mon Sep 17 00:00:00 2001
From: phix <peter.hicks@astronomer.io>
Date: Mon, 19 Feb 2024 11:09:32 -0800
Subject: [PATCH] Quality of life updates for new lineage graph display.

---
 .../src/components/ZoomPanSvg/ZoomPanSvg.tsx  |  6 +++---
 web/src/routes/column-level/ZoomControls.tsx  | 16 ++++++++++++++--
 web/src/routes/table-level/ActionBar.tsx      |  7 ++++++-
 web/src/routes/table-level/TableLevel.tsx     | 19 +++++++++++++++----
 4 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/web/libs/graph/src/components/ZoomPanSvg/ZoomPanSvg.tsx b/web/libs/graph/src/components/ZoomPanSvg/ZoomPanSvg.tsx
index c8b1fab5bf..6ed813c723 100644
--- a/web/libs/graph/src/components/ZoomPanSvg/ZoomPanSvg.tsx
+++ b/web/libs/graph/src/components/ZoomPanSvg/ZoomPanSvg.tsx
@@ -33,7 +33,7 @@ export interface ZoomPanControls {
   centerOnExtent(extent: Extent): void
   scaleZoom(kDelta?: number): void
   resetZoom(): void
-  centerOnPositionedNode(nodeId: string): void
+  centerOnPositionedNode(nodeId: string, k?: number): void
 }
 
 interface Props extends BoxProps {
@@ -232,12 +232,12 @@ export const ZoomPanSvg = ({
     animateToZoomState(zoomIdentity)
   }
 
-  const centerOnPositionedNode = (nodeId: string) => {
+  const centerOnPositionedNode = (nodeId: string, k = currentZoomState.k) => {
     const node = positionedNodes.find((node) => node.id === nodeId)
     if (!node) return
 
     const extent = getNodeExtent(node)
-    centerOnExtent(extent)
+    centerOnExtent(extent, k)
   }
 
   const fitContent = () => {
diff --git a/web/src/routes/column-level/ZoomControls.tsx b/web/src/routes/column-level/ZoomControls.tsx
index a96cbec2a8..47dd7452be 100644
--- a/web/src/routes/column-level/ZoomControls.tsx
+++ b/web/src/routes/column-level/ZoomControls.tsx
@@ -1,4 +1,4 @@
-import { CropFree, ZoomIn, ZoomOut } from '@mui/icons-material'
+import { CenterFocusStrong, CropFree, ZoomIn, ZoomOut } from '@mui/icons-material'
 import { Tooltip } from '@mui/material'
 import { theme } from '../../helpers/theme'
 import Box from '@mui/material/Box'
@@ -8,9 +8,14 @@ import React from 'react'
 interface ZoomControlsProps {
   handleScaleZoom: (inOrOut: 'in' | 'out') => void
   handleResetZoom: () => void
+  handleCenterOnNode?: () => void
 }
 
-export const ZoomControls = ({ handleScaleZoom, handleResetZoom }: ZoomControlsProps) => {
+export const ZoomControls = ({
+  handleScaleZoom,
+  handleResetZoom,
+  handleCenterOnNode,
+}: ZoomControlsProps) => {
   return (
     <Box
       display={'flex'}
@@ -38,6 +43,13 @@ export const ZoomControls = ({ handleScaleZoom, handleResetZoom }: ZoomControlsP
           <CropFree />
         </IconButton>
       </Tooltip>
+      {handleCenterOnNode && (
+        <Tooltip title={'Center on selected node'} placement={'left'}>
+          <IconButton size={'small'} onClick={handleCenterOnNode}>
+            <CenterFocusStrong />
+          </IconButton>
+        </Tooltip>
+      )}
     </Box>
   )
 }
diff --git a/web/src/routes/table-level/ActionBar.tsx b/web/src/routes/table-level/ActionBar.tsx
index 0fc7fd3a1c..25c802f623 100644
--- a/web/src/routes/table-level/ActionBar.tsx
+++ b/web/src/routes/table-level/ActionBar.tsx
@@ -109,10 +109,12 @@ export const ActionBar = ({
             control={
               <Switch
                 size={'small'}
-                defaultChecked
                 value={isFull}
+                defaultChecked={searchParams.get('isFull') === 'true'}
                 onChange={(_, checked) => {
                   setIsFull(checked)
+                  searchParams.set('isFull', checked.toString())
+                  setSearchParams(searchParams)
                 }}
               />
             }
@@ -123,8 +125,11 @@ export const ActionBar = ({
               <Switch
                 size={'small'}
                 value={isCompact}
+                defaultChecked={searchParams.get('isCompact') === 'true'}
                 onChange={(_, checked) => {
                   setIsCompact(checked)
+                  searchParams.set('isCompact', checked.toString())
+                  setSearchParams(searchParams)
                 }}
               />
             }
diff --git a/web/src/routes/table-level/TableLevel.tsx b/web/src/routes/table-level/TableLevel.tsx
index f9eb19863b..bc5f7bd7ea 100644
--- a/web/src/routes/table-level/TableLevel.tsx
+++ b/web/src/routes/table-level/TableLevel.tsx
@@ -1,8 +1,8 @@
 import * as Redux from 'redux'
 import { ActionBar } from './ActionBar'
 import { Box } from '@mui/system'
+import { DEFAULT_MAX_SCALE, Graph, ZoomPanControls } from '../../../libs/graph'
 import { Drawer } from '@mui/material'
-import { Graph, ZoomPanControls } from '../../../libs/graph'
 import { IState } from '../../store/reducers'
 import { JobOrDataset } from '../../components/lineage/types'
 import { LineageGraph } from '../../types/api'
@@ -41,8 +41,8 @@ const ColumnLevel: React.FC<ColumnLevelProps> = ({
 
   const [depth, setDepth] = useState(Number(searchParams.get('depth')) || 2)
 
-  const [isCompact, setIsCompact] = useState(false)
-  const [isFull, setIsFull] = useState(true)
+  const [isCompact, setIsCompact] = useState(searchParams.get('isCompact') === 'true')
+  const [isFull, setIsFull] = useState(searchParams.get('isFull') === 'true')
 
   const graphControls = useRef<ZoomPanControls>()
 
@@ -64,6 +64,13 @@ const ColumnLevel: React.FC<ColumnLevelProps> = ({
     graphControls.current?.fitContent()
   }
 
+  const handleCenterOnNode = () => {
+    graphControls.current?.centerOnPositionedNode(
+      `${nodeType}:${namespace}:${name}`,
+      DEFAULT_MAX_SCALE
+    )
+  }
+
   const setGraphControls = useCallbackRef((zoomControls) => {
     graphControls.current = zoomControls
   })
@@ -103,7 +110,11 @@ const ColumnLevel: React.FC<ColumnLevelProps> = ({
             <TableLevelDrawer />
           </Box>
         </Drawer>
-        <ZoomControls handleScaleZoom={handleScaleZoom} handleResetZoom={handleResetZoom} />
+        <ZoomControls
+          handleCenterOnNode={handleCenterOnNode}
+          handleScaleZoom={handleScaleZoom}
+          handleResetZoom={handleResetZoom}
+        />
         <ParentSize>
           {(parent) => (
             <Graph<JobOrDataset, TableLevelNodeData>