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

- allows a user to show/not show columns in the graph

Author: benfdking

sqlmesh/core/engine_adapter/mssql.py

@@ -6,13 +6,11 @@
 
 from sqlglot import exp
 
-from sqlmesh.core.dialect import to_schema, add_table
+from sqlmesh.core.dialect import to_schema
 from sqlmesh.core.engine_adapter.base import (
     EngineAdapterWithIndexSupport,
     EngineAdapter,
     InsertOverwriteStrategy,
-    MERGE_SOURCE_ALIAS,
-    MERGE_TARGET_ALIAS,
 )
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
@@ -34,7 +32,7 @@
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import SchemaName, TableName
-    from sqlmesh.core.engine_adapter._typing import DF, Query, QueryOrDF
+    from sqlmesh.core.engine_adapter._typing import DF, Query
 
 
 @set_catalog()
@@ -190,88 +188,7 @@ def drop_schema(
                     )
         super().drop_schema(schema_name, ignore_if_not_exists=ignore_if_not_exists, cascade=False)
 
-    def merge(
-        self,
-        target_table: TableName,
-        source_table: QueryOrDF,
-        columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
-        unique_key: t.Sequence[exp.Expression],
-        when_matched: t.Optional[exp.Whens] = None,
-        merge_filter: t.Optional[exp.Expression] = None,
-    ) -> None:
-        source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
-            source_table, columns_to_types, target_table=target_table
-        )
-        columns_to_types = columns_to_types or self.columns(target_table)
-        on = exp.and_(
-            *(
-                add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS))
-                for part in unique_key
-            )
-        )
-        if merge_filter:
-            on = exp.and_(merge_filter, on)
-
-        match_expressions = []
-        if not when_matched:
-            match_condition = None
-            unique_key_names = [y.name for y in unique_key]
-            columns_to_types_no_keys = [c for c in columns_to_types if c not in unique_key_names]
-
-            target_columns_no_keys = [
-                exp.column(c, MERGE_TARGET_ALIAS) for c in columns_to_types_no_keys
-            ]
-            source_columns_no_keys = [
-                exp.column(c, MERGE_SOURCE_ALIAS) for c in columns_to_types_no_keys
-            ]
-
-            match_condition = exp.Exists(
-                this=exp.select(*target_columns_no_keys).except_(
-                    exp.select(*source_columns_no_keys)
-                )
-            )
 
-            if target_columns_no_keys:
-                match_expressions.append(
-                    exp.When(
-                        matched=True,
-                        source=False,
-                        condition=match_condition,
-                        then=exp.Update(
-                            expressions=[
-                                exp.column(col, MERGE_TARGET_ALIAS).eq(
-                                    exp.column(col, MERGE_SOURCE_ALIAS)
-                                )
-                                for col in columns_to_types_no_keys
-                            ],
-                        ),
-                    )
-                )
-        else:
-            match_expressions.extend(when_matched.copy().expressions)
-
-        match_expressions.append(
-            exp.When(
-                matched=False,
-                source=False,
-                then=exp.Insert(
-                    this=exp.Tuple(expressions=[exp.column(col) for col in columns_to_types]),
-                    expression=exp.Tuple(
-                        expressions=[
-                            exp.column(col, MERGE_SOURCE_ALIAS) for col in columns_to_types
-                        ]
-                    ),
-                ),
-            )
-        )
-        for source_query in source_queries:
-            with source_query as query:
-                self._merge(
-                    target_table=target_table,
-                    query=query,
-                    on=on,
-                    whens=exp.Whens(expressions=match_expressions),
-                )
 
     def _convert_df_datetime(self, df: DF, columns_to_types: t.Dict[str, exp.DataType]) -> None:
         import pandas as pd

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>
+  )
+}