feat(vscode): add ability to turn on/off columns (#4826)

Author: benfdking

vscode/extension/tests/broken_project.spec.ts

@@ -137,7 +137,7 @@ test('working project, then broken through adding double model, then refixed', a
         if (activeFrame) {
           try {
             await activeFrame
-              .getByText('raw.demographics')
+              .getByText('sushi.customers')
               .waitFor({ timeout: 1000 })
             raw_demographicsCount++
           } catch {

vscode/extension/tests/lineage_settings.spec.ts

@@ -0,0 +1,66 @@
+import { test, expect } from '@playwright/test'
+import path from 'path'
+import fs from 'fs-extra'
+import os from 'os'
+import { openLineageView, SUSHI_SOURCE_PATH } from './utils'
+import { startCodeServer, stopCodeServer } from './utils_code_server'
+
+test('Settings button is visible in the lineage view', async ({ page }) => {
+  const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'))
+  await fs.copy(SUSHI_SOURCE_PATH, tempDir)
+
+  const context = await startCodeServer({
+    tempDir,
+    placeFileWithPythonInterpreter: true,
+  })
+
+  try {
+    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+
+    await page.waitForSelector('text=models')
+
+    // Click on the models folder, excluding external_models
+    await page
+      .getByRole('treeitem', { name: 'models', exact: true })
+      .locator('a')
+      .click()
+    // Open the waiters.py model
+    await page
+      .getByRole('treeitem', { name: 'waiters.py', exact: true })
+      .locator('a')
+      .click()
+    await page.waitForSelector('text=Loaded SQLMesh Context')
+
+    // Open lineage
+    await openLineageView(page)
+
+    const iframes = page.locator('iframe')
+    const iframeCount = await iframes.count()
+    let settingsCount = 0
+
+    for (let i = 0; i < iframeCount; i++) {
+      const iframe = iframes.nth(i)
+      const contentFrame = iframe.contentFrame()
+      if (contentFrame) {
+        const activeFrame = contentFrame.locator('#active-frame').contentFrame()
+        if (activeFrame) {
+          try {
+            await activeFrame
+              .getByRole('button', {
+                name: 'Settings',
+              })
+              .waitFor({ timeout: 1000 })
+            settingsCount++
+          } catch {
+            // Continue to next iframe if this one doesn't have the error
+            continue
+          }
+        }
+      }
+    }
+
+    expect(settingsCount).toBeGreaterThan(0)
+  } finally {
+    await stopCodeServer(context)
+  }
+})

vscode/react/src/components/graph/CogIcon.tsx

@@ -0,0 +1,27 @@
+import * as React from 'react'
+
+/**
+ * CogIcon as taken from https://heroicons.com/. Slightly modified to remove fill color.
+ *
+ * @param props - SVG props
+ * @returns SVG element
+ */
+export function CogIcon(props: React.SVGProps<SVGSVGElement>): JSX.Element {
+  return (
+    <svg
+      xmlns="http://www.w3.org/2000/svg"
+      aria-hidden="true"
+      className="h-6 w-6"
+      data-slot="icon"
+      viewBox="0 0 24 24"
+      {...props}
+    >
+      <path d="M17.004 10.407c.138.435-.216.842-.672.842h-3.465a.75.75 0 0 1-.65-.375l-1.732-3c-.229-.396-.053-.907.393-1.004a5.252 5.252 0 0 1 6.126 3.537ZM8.12 8.464c.307-.338.838-.235 1.066.16l1.732 3a.75.75 0 0 1 0 .75l-1.732 3c-.229.397-.76.5-1.067.161A5.23 5.23 0 0 1 6.75 12a5.23 5.23 0 0 1 1.37-3.536Zm2.758 8.666c-.447-.098-.623-.608-.394-1.004l1.733-3.002a.75.75 0 0 1 .65-.375h3.465c.457 0 .81.407.672.842a5.252 5.252 0 0 1-6.126 3.539Z" />
+      <path
+        fillRule="evenodd"
+        d="M21 12.75a.75.75 0 1 0 0-1.5h-.783a8.22 8.22 0 0 0-.237-1.357l.734-.267a.75.75 0 1 0-.513-1.41l-.735.268a8.24 8.24 0 0 0-.689-1.192l.6-.503a.75.75 0 1 0-.964-1.149l-.6.504a8.3 8.3 0 0 0-1.054-.885l.391-.678a.75.75 0 1 0-1.299-.75l-.39.676a8.188 8.188 0 0 0-1.295-.47l.136-.77a.75.75 0 0 0-1.477-.26l-.136.77a8.36 8.36 0 0 0-1.377 0l-.136-.77a.75.75 0 1 0-1.477.26l.136.77c-.448.121-.88.28-1.294.47l-.39-.676a.75.75 0 0 0-1.3.75l.392.678a8.29 8.29 0 0 0-1.054.885l-.6-.504a.75.75 0 1 0-.965 1.149l.6.503a8.243 8.243 0 0 0-.689 1.192L3.8 8.216a.75.75 0 1 0-.513 1.41l.735.267a8.222 8.222 0 0 0-.238 1.356h-.783a.75.75 0 0 0 0 1.5h.783c.042.464.122.917.238 1.356l-.735.268a.75.75 0 0 0 .513 1.41l.735-.268c.197.417.428.816.69 1.191l-.6.504a.75.75 0 0 0 .963 1.15l.601-.505c.326.323.679.62 1.054.885l-.392.68a.75.75 0 0 0 1.3.75l.39-.679c.414.192.847.35 1.294.471l-.136.77a.75.75 0 0 0 1.477.261l.137-.772a8.332 8.332 0 0 0 1.376 0l.136.772a.75.75 0 1 0 1.477-.26l-.136-.771a8.19 8.19 0 0 0 1.294-.47l.391.677a.75.75 0 0 0 1.3-.75l-.393-.679a8.29 8.29 0 0 0 1.054-.885l.601.504a.75.75 0 0 0 .964-1.15l-.6-.503c.261-.375.492-.774.69-1.191l.735.267a.75.75 0 1 0 .512-1.41l-.734-.267c.115-.439.195-.892.237-1.356h.784Zm-2.657-3.06a6.744 6.744 0 0 0-1.19-2.053 6.784 6.784 0 0 0-1.82-1.51A6.705 6.705 0 0 0 12 5.25a6.8 6.8 0 0 0-1.225.11 6.7 6.7 0 0 0-2.15.793 6.784 6.784 0 0 0-2.952 3.489.76.76 0 0 1-.036.098A6.74 6.74 0 0 0 5.251 12a6.74 6.74 0 0 0 3.366 5.842l.009.005a6.704 6.704 0 0 0 2.18.798l.022.003a6.792 6.792 0 0 0 2.368-.004 6.704 6.704 0 0 0 2.205-.811 6.785 6.785 0 0 0 1.762-1.484l.009-.01.009-.01a6.743 6.743 0 0 0 1.18-2.066c.253-.707.39-1.469.39-2.263a6.74 6.74 0 0 0-.408-2.309Z"
+        clipRule="evenodd"
+      />
+    </svg>
+  )
+}

vscode/react/src/components/graph/Graph.css

@@ -29,6 +29,7 @@ react-flow__attribution {
   box-shadow: none;
   border: var(--vscode-button-border);
   background: var(--vscode-button-background);
+  color: var(--vscode-foreground);
 }
 .react-flow__controls-button:hover {
   background: var(--vscode-button-hoverBackground);

vscode/react/src/components/graph/ModelLineage.tsx

@@ -36,6 +36,7 @@ import { Popover } from '@headlessui/react'
 import ModelLineageDetails from './ModelLineageDetails'
 import { Divider } from '@/components/divider/Divider'
 import { type ModelLineageApiLineageModelNameGet200 } from '@/api/client'
+import { SettingsControl } from '@/components/graph/SettingsControl'
 import './Graph.css'
 
 const WITH_COLUMNS_LIMIT = 30
@@ -203,6 +204,7 @@ function ModelColumnLineage(): JSX.Element {
     showControls,
     handleError,
     setActiveNodes,
+    setWithColumns,
   } = useLineageFlow()
 
   const { setCenter } = useReactFlow()
@@ -386,7 +388,12 @@ function ModelColumnLineage(): JSX.Element {
         <Controls
           className="bg-light p-1 rounded-md !border-none !shadow-lg"
           showInteractive={false}
-        />
+        >
+          <SettingsControl
+            showColumns={withColumns}
+            onWithColumnsChange={setWithColumns}
+          />
+        </Controls>
         <Background
           variant={BackgroundVariant.Cross}
           gap={32}

vscode/react/src/components/graph/ModelNode.tsx

@@ -141,7 +141,10 @@ export default function ModelNode({
   const isModelExternal = nodeType === EnumLineageNodeModelType.external
   const isModelSeed = nodeType === EnumLineageNodeModelType.seed
   const isModelUnknown = nodeType === EnumLineageNodeModelType.unknown
-  const showColumns = isArrayNotEmpty(columns) && isFalse(hasHighlightedNodes)
+  const showColumns =
+    nodeData.withColumns &&
+    isArrayNotEmpty(columns) &&
+    isFalse(hasHighlightedNodes)
   const isActiveNode =
     selectedNodes.size > 0 || activeNodes.size > 0 || withConnected
       ? isSelected ||

vscode/react/src/components/graph/SettingsControl.tsx

@@ -0,0 +1,55 @@
+import { Menu } from '@headlessui/react'
+import { CheckIcon } from '@heroicons/react/24/outline'
+import { CogIcon } from '@/components/graph/CogIcon'
+import clsx from 'clsx'
+
+interface SettingsControlProps {
+  showColumns: boolean
+  onWithColumnsChange: (value: boolean) => void
+}
+
+export function SettingsControl({
+  showColumns,
+  onWithColumnsChange,
+}: SettingsControlProps): JSX.Element {
+  return (
+    <Menu
+      as="div"
+      className="relative"
+    >
+      <Menu.Button
+        className="react-flow__controls-button"
+        title="Settings"
+      >
+        <CogIcon
+          className="h-3 w-3"
+          aria-hidden="true"
+        />
+      </Menu.Button>
+      <Menu.Items className="absolute bottom-0 left-full ml-2 w-56 origin-bottom-left divide-y bg-theme shadow-lg focus:outline-none z-50">
+        <Menu.Item>
+          {({ active }) => (
+            <button
+              className={clsx(
+                'group flex w-full items-center px-2 py-1 text-sm',
+                'text-[var(--vscode-button-foreground)]',
+                active
+                  ? 'bg-[var(--vscode-button-background)]'
+                  : 'bg-[var(--vscode-button-hoverBackground)]',
+              )}
+              onClick={() => onWithColumnsChange(!showColumns)}
+            >
+              <span className="flex-1 text-left">Show Columns</span>
+              {showColumns && (
+                <CheckIcon
+                  className="h-4 w-4 text-primary-500"
+                  aria-hidden="true"
+                />
+              )}
+            </button>
+          )}
+        </Menu.Item>
+      </Menu.Items>
+    </Menu>
+  )
+}