Fix: Re-fetch model's lineage after change to models/columns in file (#891)

Author: mykalmax

web/client/src/library/components/graph/Graph.tsx

@@ -679,6 +679,8 @@ function ModelColumnLineage({
           onNodesChange={onNodesChange}
           onEdgesChange={onEdgesChange}
           nodeOrigin={[0.5, 0.5]}
+          minZoom={0.1}
+          maxZoom={1.5}
           fitView
         >
           <Controls className="bg-light p-1 rounded-md !border-none !shadow-lg" />

web/client/src/library/components/graph/ModelLineage.tsx

@@ -17,15 +17,15 @@ export default function ModelLineage({
   highlightedNodes?: Record<string, string[]>
   className?: string
 }): JSX.Element {
-  const { clearActiveEdges, setLineage } = useLineageFlow()
+  const { clearActiveEdges, setLineage, models } = useLineageFlow()
 
   const { data: dataLineage, refetch: getModelLineage } = useApiModelLineage(
     model.name,
   )
 
   const debouncedGetModelLineage = useCallback(
     debounceAsync(getModelLineage, 500),
-    [model.name],
+    [model.name, models],
   )
 
   useEffect(() => {

web/client/src/library/components/graph/help.ts

@@ -177,11 +177,14 @@ function getNodeMap({
   const CHAR_WIDTH = 8
   const MAX_VISIBLE_COLUMNS = 5
 
-  const current = nodes.reduce(
-    (acc: Record<string, Node>, node) =>
-      Object.assign(acc, { [node.id]: node }),
-    {},
-  )
+  const current = nodes.reduce((acc: Record<string, Node>, node) => {
+    // Checking if any nodes have been removed from the graph
+    if (models.has(node.id)) {
+      acc[node.id] = node
+    }
+
+    return acc
+  }, {})
 
   return modelNames.reduce((acc: Record<string, Node>, label: string) => {
     const node =

web/client/src/library/pages/ide/IDE.tsx

@@ -15,11 +15,11 @@ import { EnumSize, EnumVariant } from '~/types/enum'
 import { Link, Outlet, useLocation, useNavigate } from 'react-router-dom'
 import { EnumRoutes } from '~/routes'
 import { useStoreFileTree } from '@context/fileTree'
-import { Divider } from '@components/divider/Divider'
+import IDEProvider, { EnumErrorKey, useIDE } from './context'
+import { type Model } from '@api/client'
 import { Button } from '@components/button/Button'
+import { Divider } from '@components/divider/Divider'
 import Container from '@components/container/Container'
-import { EnumErrorKey, useIDE } from './context'
-import { type Model } from '@api/client'
 
 const ReportErrors = lazy(
   async () => await import('../../components/report/ReportErrors'),
@@ -114,35 +114,39 @@ export default function PageIDE(): JSX.Element {
 
   return (
     <Container.Page>
-      <PlanSidebar />
-      <div className="w-full flex justify-between items-center min-h-[2rem] z-50">
-        <div className="px-3 flex items-center whitespace-nowrap">
-          <h3 className="font-bold text-primary-500">
-            <span className="inline-block">/</span>
-            {project?.name}
-          </h3>
-          <ArrowLongRightIcon className="w-8 mx-4 text-neutral-50" />
-          <Button
-            size={EnumSize.sm}
-            variant={EnumVariant.Neutral}
-          >
-            {isActivePageEditor ? (
-              <Link to={EnumRoutes.IdeDocs}>Docs</Link>
-            ) : (
-              <Link to={EnumRoutes.IdeEditor}>Editor</Link>
-            )}
-          </Button>
-        </div>
-        <div className="px-3 flex items-center min-w-[10rem] justify-end">
-          <RunPlan />
-          <Suspense>
-            {activePlan != null && <ActivePlan plan={activePlan} />}
-          </Suspense>
-          <ReportErrors />
+      <IDEProvider>
+        <PlanSidebar />
+        <div className="w-full flex justify-between items-center min-h-[2rem] z-50">
+          <div className="px-3 flex items-center whitespace-nowrap">
+            <h3 className="font-bold text-primary-500">
+              <span className="inline-block">/</span>
+              {project?.name}
+            </h3>
+            <ArrowLongRightIcon className="w-8 mx-4 text-neutral-50" />
+            <Button
+              size={EnumSize.sm}
+              variant={EnumVariant.Neutral}
+            >
+              {isActivePageEditor ? (
+                <Link to={EnumRoutes.IdeDocs}>Docs</Link>
+              ) : (
+                <Link to={EnumRoutes.IdeEditor}>Editor</Link>
+              )}
+            </Button>
+          </div>
+          {isActivePageEditor && (
+            <div className="px-3 flex items-center min-w-[10rem] justify-end">
+              <RunPlan />
+              <Suspense>
+                {activePlan != null && <ActivePlan plan={activePlan} />}
+              </Suspense>
+              <ReportErrors />
+            </div>
+          )}
         </div>
-      </div>
-      <Divider />
-      <Outlet />
+        <Divider />
+        <Outlet />
+      </IDEProvider>
     </Container.Page>
   )
 }

web/client/src/routes.tsx

@@ -1,5 +1,4 @@
 import { createBrowserRouter } from 'react-router-dom'
-import IDEProvider from './library/pages/ide/context'
 import { Suspense, lazy } from 'react'
 import NotFound from './library/pages/root/NotFound'
 import Loading from '@components/loading/Loading'
@@ -25,11 +24,7 @@ export const EnumRoutes = {
 export const router = createBrowserRouter([
   {
     path: '/',
-    element: (
-      <IDEProvider>
-        <IDE />
-      </IDEProvider>
-    ),
+    element: <IDE />,
     children: [
       {
         path: 'editor',

web/server/api/endpoints/plan.py

@@ -72,7 +72,7 @@ async def run_plan(
                 if interval.snapshot_name in plan.context_diff.snapshots
                 else interval.snapshot_name,
                 interval=[
-                    tuple([to_ds(t) for t in make_inclusive(start, end)])
+                    tuple(to_ds(t) for t in make_inclusive(start, end))
                     for start, end in interval.merged_intervals
                 ][0],
                 batches=tasks.get(interval.snapshot_name, 0),