From 40254df2628ea4d77bb4fdb63c0912ffb1be3edd Mon Sep 17 00:00:00 2001 From: Max Mykal Date: Tue, 7 Oct 2025 15:45:30 -0700 Subject: [PATCH 1/2] add more generic types and fix satories --- pnpm-lock.yaml | 45 ++-- web/common/package.json | 3 - .../ColumnLevelLineageContext.ts | 91 +++---- .../LineageColumnLevel/FactoryColumn.tsx | 21 +- .../Lineage/LineageColumnLevel/help.ts | 78 +++--- .../useColumnLevelLineage.ts | 23 +- .../src/components/Lineage/LineageContext.ts | 62 ++++- .../src/components/Lineage/LineageLayout.tsx | 10 +- .../components/Lineage/LineageLayoutBase.tsx | 87 +++++-- .../Lineage/edge/FactoryEdgeWithGradient.tsx | 10 +- web/common/src/components/Lineage/help.ts | 61 +++-- .../components/Lineage/layout/dagreLayout.ts | 15 +- .../src/components/Lineage/layout/help.ts | 43 +++- .../components/Lineage/node/NodeHandle.tsx | 7 +- .../components/Lineage/node/NodeHandles.tsx | 16 +- .../src/components/Lineage/node/NodePort.tsx | 24 +- .../components/Lineage/node/base-handle.tsx | 1 - .../Lineage/stories/Lineage.stories.tsx | 228 +++++++++--------- .../Lineage/stories/ModelLineage.tsx | 171 +++++++------ .../Lineage/stories/ModelLineageContext.ts | 64 ++++- .../Lineage/stories/ModelNodeColumn.tsx | 11 +- web/common/src/components/Lineage/utils.ts | 54 +++-- web/common/src/index.ts | 2 + web/common/src/types.ts | 40 ++- 24 files changed, 745 insertions(+), 422 deletions(-) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 2fec93a8f3..aeacb362d0 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -435,9 +435,6 @@ importers: '@types/dagre': specifier: 0.7.53 version: 0.7.53 - '@types/lodash': - specifier: 4.17.20 - version: 4.17.20 '@types/node': specifier: 20.11.25 version: 20.11.25 @@ -495,9 +492,6 @@ importers: globals: specifier: 16.3.0 version: 16.3.0 - lodash: - specifier: 4.17.21 - version: 4.17.21 lucide-react: specifier: 0.542.0 version: 0.542.0(react@18.3.1) @@ -2662,9 +2656,6 @@ packages: '@types/jsonfile@6.1.4': resolution: {integrity: sha512-D5qGUYwjvnNNextdU59/+fI+spnwtTFmyQP0h+PfIOSkNfpU6AOICUOkm4i0OnSk+NyjdPJrxCDro0sJsWlRpQ==} - '@types/lodash@4.17.20': - resolution: {integrity: sha512-H3MHACvFUEiujabxhaI/ImO6gUrd8oOurg7LQtS7mbwIXA/cUqWrvBsaeJ23aZEPk1TAYkurjfMbSELfoCXlGA==} - '@types/mdast@4.0.4': resolution: {integrity: sha512-kGaNbPh1k7AFzgpud/gMdvIm5xuECykRR+JnWKQno9TAXVa6WIVCGTPvYGekIDL4uwCZQSYbUxNBSb1aUo79oA==} @@ -7064,7 +7055,7 @@ snapshots: '@babel/traverse': 7.28.0 '@babel/types': 7.28.1 convert-source-map: 2.0.0 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 gensync: 1.0.0-beta.2 json5: 2.2.3 semver: 6.3.1 @@ -7232,7 +7223,7 @@ snapshots: '@babel/parser': 7.28.0 '@babel/template': 7.27.2 '@babel/types': 7.28.1 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 transitivePeerDependencies: - supports-color @@ -7489,7 +7480,7 @@ snapshots: '@eslint/config-array@0.21.0': dependencies: '@eslint/object-schema': 2.1.6 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 minimatch: 3.1.2 transitivePeerDependencies: - supports-color @@ -7503,7 +7494,7 @@ snapshots: '@eslint/eslintrc@3.3.1': dependencies: ajv: 6.12.6 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 espree: 10.4.0 globals: 14.0.0 ignore: 5.3.2 @@ -9425,8 +9416,6 @@ snapshots: dependencies: '@types/node': 20.11.25 - '@types/lodash@4.17.20': {} - '@types/mdast@4.0.4': dependencies: '@types/unist': 3.0.3 @@ -9502,7 +9491,7 @@ snapshots: '@typescript-eslint/types': 8.38.0 '@typescript-eslint/typescript-estree': 8.38.0(typescript@5.8.3) '@typescript-eslint/visitor-keys': 8.38.0 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 eslint: 9.31.0(jiti@2.4.2) typescript: 5.8.3 transitivePeerDependencies: @@ -9512,7 +9501,7 @@ snapshots: dependencies: '@typescript-eslint/tsconfig-utils': 8.38.0(typescript@5.8.3) '@typescript-eslint/types': 8.38.0 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 typescript: 5.8.3 transitivePeerDependencies: - supports-color @@ -9531,7 +9520,7 @@ snapshots: '@typescript-eslint/types': 8.38.0 '@typescript-eslint/typescript-estree': 8.38.0(typescript@5.8.3) '@typescript-eslint/utils': 8.38.0(eslint@9.31.0(jiti@2.4.2))(typescript@5.8.3) - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 eslint: 9.31.0(jiti@2.4.2) ts-api-utils: 2.1.0(typescript@5.8.3) typescript: 5.8.3 @@ -9546,7 +9535,7 @@ snapshots: '@typescript-eslint/tsconfig-utils': 8.38.0(typescript@5.8.3) '@typescript-eslint/types': 8.38.0 '@typescript-eslint/visitor-keys': 8.38.0 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 fast-glob: 3.3.3 is-glob: 4.0.3 minimatch: 9.0.5 @@ -10644,6 +10633,10 @@ snapshots: de-indent@1.0.2: {} + debug@4.4.1: + dependencies: + ms: 2.1.3 + debug@4.4.1(supports-color@8.1.1): dependencies: ms: 2.1.3 @@ -10916,7 +10909,7 @@ snapshots: esbuild-register@3.6.0(esbuild@0.25.8): dependencies: - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 esbuild: 0.25.8 transitivePeerDependencies: - supports-color @@ -10999,7 +10992,7 @@ snapshots: ajv: 6.12.6 chalk: 4.1.2 cross-spawn: 7.0.6 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 escape-string-regexp: 4.0.0 eslint-scope: 8.4.0 eslint-visitor-keys: 4.2.1 @@ -11410,7 +11403,7 @@ snapshots: http-proxy-agent@7.0.2: dependencies: agent-base: 7.1.4 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 transitivePeerDependencies: - supports-color @@ -11419,7 +11412,7 @@ snapshots: https-proxy-agent@7.0.6: dependencies: agent-base: 7.1.4 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 transitivePeerDependencies: - supports-color @@ -13992,7 +13985,7 @@ snapshots: vite-node@3.2.4(@types/node@20.11.25)(jiti@2.4.2)(lightningcss@1.30.1)(terser@5.44.0)(tsx@4.20.3)(yaml@2.8.0): dependencies: cac: 6.7.14 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 es-module-lexer: 1.7.0 pathe: 2.0.3 vite: 6.3.5(@types/node@20.11.25)(jiti@2.4.2)(lightningcss@1.30.1)(terser@5.44.0)(tsx@4.20.3)(yaml@2.8.0) @@ -14042,7 +14035,7 @@ snapshots: '@volar/typescript': 2.4.23 '@vue/language-core': 2.2.0(typescript@5.8.3) compare-versions: 6.1.1 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 kolorist: 1.8.0 local-pkg: 1.1.1 magic-string: 0.30.17 @@ -14108,7 +14101,7 @@ snapshots: '@vitest/spy': 3.2.4 '@vitest/utils': 3.2.4 chai: 5.2.1 - debug: 4.4.1(supports-color@8.1.1) + debug: 4.4.1 expect-type: 1.2.2 magic-string: 0.30.17 pathe: 2.0.3 diff --git a/web/common/package.json b/web/common/package.json index ef91337174..6a0965f19e 100644 --- a/web/common/package.json +++ b/web/common/package.json @@ -14,7 +14,6 @@ "@testing-library/react": "16.3.0", "@testing-library/user-event": "14.6.1", "@types/dagre": "0.7.53", - "@types/lodash": "4.17.20", "@types/node": "20.11.25", "@types/react": "18.3.23", "@types/react-dom": "18.3.7", @@ -34,7 +33,6 @@ "eslint-plugin-storybook": "9.1.5", "fuse.js": "7.1.0", "globals": "16.3.0", - "lodash": "4.17.21", "lucide-react": "0.542.0", "playwright": "1.54.1", "postcss": "8.5.6", @@ -95,7 +93,6 @@ "dagre": "0.8.5", "deepmerge": "4.3.1", "fuse.js": "7.1.0", - "lodash": "4.17.21", "lucide-react": "0.542.0", "react": "18.3.1", "react-dom": "18.3.1", diff --git a/web/common/src/components/Lineage/LineageColumnLevel/ColumnLevelLineageContext.ts b/web/common/src/components/Lineage/LineageColumnLevel/ColumnLevelLineageContext.ts index 227fc70394..4dd6ca93ef 100644 --- a/web/common/src/components/Lineage/LineageColumnLevel/ColumnLevelLineageContext.ts +++ b/web/common/src/components/Lineage/LineageColumnLevel/ColumnLevelLineageContext.ts @@ -2,64 +2,36 @@ import React from 'react' import { type PortId } from '../utils' -export type LineageColumn = { - source?: string | null - expression?: string | null - models: Record -} - -export type ColumnLevelModelConnections< - TAdjacencyListKey extends string, - TAdjacencyListColumnKey extends string, -> = Record -export type ColumnLevelDetails< - TAdjacencyListKey extends string, - TAdjacencyListColumnKey extends string, -> = Omit & { - models: ColumnLevelModelConnections< - TAdjacencyListKey, - TAdjacencyListColumnKey - > -} -export type ColumnLevelConnections< - TAdjacencyListKey extends string, - TAdjacencyListColumnKey extends string, -> = Record< - TAdjacencyListColumnKey, - ColumnLevelDetails -> export type ColumnLevelLineageAdjacencyList< TAdjacencyListKey extends string, TAdjacencyListColumnKey extends string, -> = Record< - TAdjacencyListKey, - ColumnLevelConnections -> +> = { + [K in TAdjacencyListKey]: { + [C in TAdjacencyListColumnKey]: { + source?: string | null + expression?: string | null + models: Record + } + } +} export type ColumnLevelLineageContextValue< TAdjacencyListKey extends string, TAdjacencyListColumnKey extends string, TColumnID extends string = PortId, -> = { - adjacencyListColumnLevel: ColumnLevelLineageAdjacencyList< + TColumnLevelLineageAdjacencyList extends ColumnLevelLineageAdjacencyList< TAdjacencyListKey, TAdjacencyListColumnKey - > + > = ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + >, +> = { + adjacencyListColumnLevel: TColumnLevelLineageAdjacencyList selectedColumns: Set - columnLevelLineage: Map< - TColumnID, - ColumnLevelLineageAdjacencyList - > + columnLevelLineage: Map setColumnLevelLineage: React.Dispatch< - React.SetStateAction< - Map< - TColumnID, - ColumnLevelLineageAdjacencyList< - TAdjacencyListKey, - TAdjacencyListColumnKey - > - > - > + React.SetStateAction> > showColumns: boolean setShowColumns: React.Dispatch> @@ -71,16 +43,17 @@ export function getColumnLevelLineageContextInitial< TAdjacencyListKey extends string, TAdjacencyListColumnKey extends string, TColumnID extends string = PortId, + TColumnLevelLineageAdjacencyList extends ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + > = ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + >, >() { return { - adjacencyListColumnLevel: {}, - columnLevelLineage: new Map< - TColumnID, - ColumnLevelLineageAdjacencyList< - TAdjacencyListKey, - TAdjacencyListColumnKey - > - >(), + adjacencyListColumnLevel: {} as TColumnLevelLineageAdjacencyList, + columnLevelLineage: new Map(), setColumnLevelLineage: () => {}, showColumns: false, setShowColumns: () => {}, @@ -94,8 +67,16 @@ export type ColumnLevelLineageContextHook< TAdjacencyListKey extends string, TAdjacencyListColumnKey extends string, TColumnID extends string = PortId, + TColumnLevelLineageAdjacencyList extends ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + > = ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + >, > = () => ColumnLevelLineageContextValue< TAdjacencyListKey, TAdjacencyListColumnKey, - TColumnID + TColumnID, + TColumnLevelLineageAdjacencyList > diff --git a/web/common/src/components/Lineage/LineageColumnLevel/FactoryColumn.tsx b/web/common/src/components/Lineage/LineageColumnLevel/FactoryColumn.tsx index 90def0f5ea..350437c16e 100644 --- a/web/common/src/components/Lineage/LineageColumnLevel/FactoryColumn.tsx +++ b/web/common/src/components/Lineage/LineageColumnLevel/FactoryColumn.tsx @@ -10,7 +10,7 @@ import React from 'react' import { cn } from '@/utils' import { NodeBadge } from '../node/NodeBadge' import { NodePort } from '../node/NodePort' -import { type NodeId, type PortId } from '../utils' +import { type NodeId, type PortHandleId, type PortId } from '../utils' import { type ColumnLevelLineageAdjacencyList, type ColumnLevelLineageContextHook, @@ -28,11 +28,21 @@ export function FactoryColumn< TAdjacencyListColumnKey extends string, TNodeID extends string = NodeId, TColumnID extends string = PortId, + TLeftPortHandleId extends string = PortHandleId, + TRightPortHandleId extends string = PortHandleId, + TColumnLevelLineageAdjacencyList extends ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + > = ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + >, >( useLineage: ColumnLevelLineageContextHook< TAdjacencyListKey, TAdjacencyListColumnKey, - TColumnID + TColumnID, + TColumnLevelLineageAdjacencyList >, ) { return React.memo(function FactoryColumn({ @@ -59,10 +69,7 @@ export function FactoryColumn< type: string description?: string | null className?: string - data?: ColumnLevelLineageAdjacencyList< - TAdjacencyListKey, - TAdjacencyListColumnKey - > + data?: TColumnLevelLineageAdjacencyList isFetching?: boolean error?: Error | null renderError?: (error: Error) => React.ReactNode @@ -248,7 +255,7 @@ export function FactoryColumn< } return isSelectedColumn ? ( - id={id} nodeId={nodeId} className={cn( diff --git a/web/common/src/components/Lineage/LineageColumnLevel/help.ts b/web/common/src/components/Lineage/LineageColumnLevel/help.ts index 30115450cd..fe75ed162a 100644 --- a/web/common/src/components/Lineage/LineageColumnLevel/help.ts +++ b/web/common/src/components/Lineage/LineageColumnLevel/help.ts @@ -9,28 +9,26 @@ import { type PortId, type TransformEdgeFn, } from '../utils' -import { - type ColumnLevelConnections, - type ColumnLevelDetails, - type ColumnLevelLineageAdjacencyList, -} from './ColumnLevelLineageContext' +import { type ColumnLevelLineageAdjacencyList } from './ColumnLevelLineageContext' export const MAX_COLUMNS_TO_DISPLAY = 5 export function getAdjacencyListKeysFromColumnLineage< TAdjacencyListKey extends string, TAdjacencyListColumnKey extends string, ->( - columnLineage: ColumnLevelLineageAdjacencyList< + TColumnLevelLineageAdjacencyList extends ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + > = ColumnLevelLineageAdjacencyList< TAdjacencyListKey, TAdjacencyListColumnKey >, -) { +>(columnLineage: TColumnLevelLineageAdjacencyList) { const adjacencyListKeys = new Set() const targets = Object.entries(columnLineage) as [ TAdjacencyListKey, - ColumnLevelConnections, + TColumnLevelLineageAdjacencyList[TAdjacencyListKey], ][] for (const [sourceModelName, targetColumns] of targets) { @@ -38,7 +36,7 @@ export function getAdjacencyListKeysFromColumnLineage< const targetConnections = Object.entries(targetColumns) as [ TAdjacencyListColumnKey, - ColumnLevelDetails, + TColumnLevelLineageAdjacencyList[TAdjacencyListKey][TAdjacencyListColumnKey], ][] for (const [, { models: sourceModels }] of targetConnections) { @@ -58,32 +56,52 @@ export function getEdgesFromColumnLineage< TAdjacencyListColumnKey extends string, TEdgeData extends LineageEdgeData = LineageEdgeData, TEdgeID extends string = EdgeId, - TNodeID extends string = NodeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, + TColumnLevelLineageAdjacencyList extends ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + > = ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + >, >({ columnLineage, transformEdge, }: { - columnLineage: ColumnLevelLineageAdjacencyList< - TAdjacencyListKey, - TAdjacencyListColumnKey + columnLineage: TColumnLevelLineageAdjacencyList + transformEdge: TransformEdgeFn< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID > - transformEdge: TransformEdgeFn }) { - const edges: LineageEdge[] = [] - const modelLevelEdgeIDs = new Map() + const edges: LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >[] = [] + const modelLevelEdgeIDs = new Map() const targets = Object.entries(columnLineage || {}) as [ TAdjacencyListKey, - ColumnLevelConnections, + TColumnLevelLineageAdjacencyList[TAdjacencyListKey], ][] for (const [targetModelName, targetColumns] of targets) { const targetConnections = Object.entries(targetColumns) as [ TAdjacencyListColumnKey, - ColumnLevelDetails, + TColumnLevelLineageAdjacencyList[TAdjacencyListKey][TAdjacencyListColumnKey], ][] - const targetNodeId = toNodeID(targetModelName) + const targetNodeId = toNodeID(targetModelName) for (const [ targetColumnName, @@ -95,7 +113,7 @@ export function getEdgesFromColumnLineage< ][] for (const [sourceModelName, sourceColumns] of sources) { - const sourceNodeId = toNodeID(sourceModelName) + const sourceNodeId = toNodeID(sourceModelName) modelLevelEdgeIDs.set( toEdgeID(sourceModelName, targetModelName), @@ -109,11 +127,11 @@ export function getEdgesFromColumnLineage< targetModelName, targetColumnName, ) - const sourceColumnId = toPortID( + const sourceColumnId = toPortID( sourceModelName, sourceColumnName, ) - const targetColumnId = toPortID( + const targetColumnId = toPortID( targetModelName, targetColumnName, ) @@ -145,22 +163,24 @@ export function getConnectedColumnsIDs< TAdjacencyListKey extends string, TAdjacencyListColumnKey extends string, TColumnID extends string = PortId, ->( - adjacencyList: ColumnLevelLineageAdjacencyList< + TColumnLevelLineageAdjacencyList extends ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + > = ColumnLevelLineageAdjacencyList< TAdjacencyListKey, TAdjacencyListColumnKey >, -) { +>(adjacencyList: TColumnLevelLineageAdjacencyList) { const connectedColumns = new Set() const targets = Object.entries(adjacencyList) as [ TAdjacencyListKey, - ColumnLevelConnections, + TColumnLevelLineageAdjacencyList[TAdjacencyListKey], ][] for (const [sourceModelName, targetColumns] of targets) { const targetConnections = Object.entries(targetColumns) as [ TAdjacencyListColumnKey, - ColumnLevelDetails, + TColumnLevelLineageAdjacencyList[TAdjacencyListKey][TAdjacencyListColumnKey], ][] for (const [ diff --git a/web/common/src/components/Lineage/LineageColumnLevel/useColumnLevelLineage.ts b/web/common/src/components/Lineage/LineageColumnLevel/useColumnLevelLineage.ts index da1a6b8ee8..53032c2c12 100644 --- a/web/common/src/components/Lineage/LineageColumnLevel/useColumnLevelLineage.ts +++ b/web/common/src/components/Lineage/LineageColumnLevel/useColumnLevelLineage.ts @@ -12,19 +12,18 @@ export function useColumnLevelLineage< TAdjacencyListKey extends string, TAdjacencyListColumnKey extends string, TColumnID extends string = PortId, ->( - columnLevelLineage: Map< - TColumnID, - ColumnLevelLineageAdjacencyList + TColumnLevelLineageAdjacencyList extends ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey + > = ColumnLevelLineageAdjacencyList< + TAdjacencyListKey, + TAdjacencyListColumnKey >, -) { +>(columnLevelLineage: Map) { const adjacencyListColumnLevel = React.useMemo(() => { return merge.all(Array.from(columnLevelLineage.values()), { arrayMerge: (dest, source) => Array.from(new Set([...dest, ...source])), - }) as ColumnLevelLineageAdjacencyList< - TAdjacencyListKey, - TAdjacencyListColumnKey - > + }) as TColumnLevelLineageAdjacencyList }, [columnLevelLineage]) const selectedColumns = React.useMemo(() => { @@ -37,7 +36,11 @@ export function useColumnLevelLineage< const adjacencyListKeysColumnLevel = React.useMemo(() => { return adjacencyListColumnLevel != null - ? getAdjacencyListKeysFromColumnLineage(adjacencyListColumnLevel) + ? getAdjacencyListKeysFromColumnLineage< + TAdjacencyListKey, + TAdjacencyListColumnKey, + TColumnLevelLineageAdjacencyList + >(adjacencyListColumnLevel) : [] }, [adjacencyListColumnLevel]) diff --git a/web/common/src/components/Lineage/LineageContext.ts b/web/common/src/components/Lineage/LineageContext.ts index 9da54dcbee..4a90031217 100644 --- a/web/common/src/components/Lineage/LineageContext.ts +++ b/web/common/src/components/Lineage/LineageContext.ts @@ -17,7 +17,10 @@ export interface LineageContextValue< TEdgeData extends LineageEdgeData = LineageEdgeData, TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = TNodeID, + TTargetID extends string = TNodeID, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, > { // Node selection showOnlySelectedNodes: boolean @@ -34,9 +37,25 @@ export interface LineageContextValue< setZoom: React.Dispatch> // Nodes and Edges - edges: LineageEdge[] + edges: LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >[] setEdges: React.Dispatch< - React.SetStateAction[]> + React.SetStateAction< + LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >[] + > > nodes: LineageNode[] nodesMap: LineageNodesMap @@ -73,22 +92,49 @@ export type LineageContextHook< TEdgeData extends LineageEdgeData = LineageEdgeData, TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, -> = () => LineageContextValue + TSourceID extends string = TNodeID, + TTargetID extends string = TNodeID, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, +> = () => LineageContextValue< + TNodeData, + TEdgeData, + TNodeID, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID +> export function createLineageContext< TNodeData extends LineageNodeData = LineageNodeData, TEdgeData extends LineageEdgeData = LineageEdgeData, TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = TNodeID, + TTargetID extends string = TNodeID, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, TLineageContextValue extends LineageContextValue< TNodeData, TEdgeData, TNodeID, TEdgeID, - TPortID - > = LineageContextValue, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + > = LineageContextValue< + TNodeData, + TEdgeData, + TNodeID, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >, >(initial: TLineageContextValue) { const LineageContext = React.createContext(initial) diff --git a/web/common/src/components/Lineage/LineageLayout.tsx b/web/common/src/components/Lineage/LineageLayout.tsx index 2ab4a34879..a9b5ec512f 100644 --- a/web/common/src/components/Lineage/LineageLayout.tsx +++ b/web/common/src/components/Lineage/LineageLayout.tsx @@ -26,7 +26,10 @@ export function LineageLayout< TEdgeData extends LineageEdgeData = LineageEdgeData, TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = TNodeID, + TTargetID extends string = TNodeID, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, >({ nodeTypes, edgeTypes, @@ -44,7 +47,10 @@ export function LineageLayout< TEdgeData, TNodeID, TEdgeID, - TPortID + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID > isBuildingLayout?: boolean nodeTypes?: NodeTypes diff --git a/web/common/src/components/Lineage/LineageLayoutBase.tsx b/web/common/src/components/Lineage/LineageLayoutBase.tsx index a21c1bac17..6d3975d19a 100644 --- a/web/common/src/components/Lineage/LineageLayoutBase.tsx +++ b/web/common/src/components/Lineage/LineageLayoutBase.tsx @@ -20,7 +20,6 @@ import { import '@xyflow/react/dist/style.css' import './Lineage.css' -import { debounce } from 'lodash' import { CircuitBoard, Crosshair, LocateFixed, RotateCcw } from 'lucide-react' import React from 'react' @@ -39,8 +38,8 @@ import { NODES_TRESHOLD_ZOOM, type NodeId, type EdgeId, - ZOOM_THRESHOLD, type PortId, + ZOOM_THRESHOLD, } from './utils' import '@xyflow/react/dist/style.css' @@ -50,9 +49,12 @@ import { cn } from '@/utils' export function LineageLayoutBase< TNodeData extends LineageNodeData = LineageNodeData, TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TNodeID extends string = NodeId, + TSourceID extends string = TNodeID, + TTargetID extends string = TNodeID, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, >({ nodeTypes, edgeTypes, @@ -69,7 +71,10 @@ export function LineageLayoutBase< TEdgeData, TNodeID, TEdgeID, - TPortID + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID > nodesDraggable?: boolean nodesConnectable?: boolean @@ -106,8 +111,19 @@ export function LineageLayoutBase< setSelectedEdges, } = useLineage() - const [nodes, setNodes] = React.useState(initialNodes) - const [edges, setEdges] = React.useState(initialEdges) + const [nodes, setNodes] = React.useState[]>( + [], + ) + const [edges, setEdges] = React.useState< + LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >[] + >([]) const onNodesChange = React.useCallback( (changes: NodeChange>[]) => { @@ -120,13 +136,28 @@ export function LineageLayoutBase< const onEdgesChange = React.useCallback( ( - changes: EdgeChange>[], + changes: EdgeChange< + LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + > + >[], ) => { setEdges( - applyEdgeChanges>( - changes, - edges, - ), + applyEdgeChanges< + LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + > + >(changes, edges), ) }, [edges, setEdges], @@ -235,12 +266,23 @@ export function LineageLayoutBase< const connectedEdges = getConnectedEdges< LineageNode, - LineageEdge + LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + > >(connectedNodes, edges) const selectedNodes = new Set(connectedNodes.map(node => node.id)) const selectedEdges = new Set( connectedEdges.reduce((acc, edge) => { - if ([edge.source, edge.target].every(id => selectedNodes.has(id))) { + if ( + [edge.source, edge.target].every(id => + selectedNodes.has(id as unknown as TNodeID), + ) + ) { edge.zIndex = 2 acc.add(edge.id) } else { @@ -278,7 +320,14 @@ export function LineageLayoutBase< return ( , - LineageEdge + LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + > > className={cn('shrink-0', className)} nodes={nodes} @@ -351,3 +400,11 @@ export function LineageLayoutBase< ) } + +function debounce(func: T, wait: number) { + let timeout: NodeJS.Timeout + return (...args: unknown[]) => { + clearTimeout(timeout) + timeout = setTimeout(() => func(...args), wait) + } +} diff --git a/web/common/src/components/Lineage/edge/FactoryEdgeWithGradient.tsx b/web/common/src/components/Lineage/edge/FactoryEdgeWithGradient.tsx index a89027ffef..aee8790b35 100644 --- a/web/common/src/components/Lineage/edge/FactoryEdgeWithGradient.tsx +++ b/web/common/src/components/Lineage/edge/FactoryEdgeWithGradient.tsx @@ -15,14 +15,20 @@ export function FactoryEdgeWithGradient< TEdgeData extends EdgeData = EdgeData, TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = TNodeID, + TTargetID extends string = TNodeID, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, >( useLineage: LineageContextHook< TNodeData, TEdgeData, TNodeID, TEdgeID, - TPortID + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID >, ) { return React.memo(({ data, id, ...props }: EdgeProps>) => { diff --git a/web/common/src/components/Lineage/help.ts b/web/common/src/components/Lineage/help.ts index 97f4ad9542..e8041d9f56 100644 --- a/web/common/src/components/Lineage/help.ts +++ b/web/common/src/components/Lineage/help.ts @@ -60,13 +60,22 @@ export function getTransformedNodes< export function getTransformedModelEdgesSourceTargets< TAdjacencyListKey extends string, TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, >( adjacencyListKeys: TAdjacencyListKey[], lineageAdjacencyList: LineageAdjacencyList, - transformEdge: TransformEdgeFn, + transformEdge: TransformEdgeFn< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >, ) { const nodesCount = adjacencyListKeys.length @@ -76,7 +85,7 @@ export function getTransformedModelEdgesSourceTargets< for (let i = 0; i < nodesCount; i++) { const sourceAdjacencyListKey = adjacencyListKeys[i] - const sourceNodeId = toNodeID(sourceAdjacencyListKey) + const sourceNodeId = toNodeID(sourceAdjacencyListKey) const targets = lineageAdjacencyList[sourceAdjacencyListKey] const targetsCount = targets?.length || 0 @@ -91,7 +100,7 @@ export function getTransformedModelEdgesSourceTargets< sourceAdjacencyListKey, targetAdjacencyListKey, ) - const targetNodeId = toNodeID(targetAdjacencyListKey) + const targetNodeId = toNodeID(targetAdjacencyListKey) edges.push(transformEdge('edge', edgeId, sourceNodeId, targetNodeId)) } @@ -103,13 +112,22 @@ export function getTransformedModelEdgesSourceTargets< export function getTransformedModelEdgesTargetSources< TAdjacencyListKey extends string, TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, >( adjacencyListKeys: TAdjacencyListKey[], lineageAdjacencyList: LineageAdjacencyList, - transformEdge: TransformEdgeFn, + transformEdge: TransformEdgeFn< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >, ) { const nodesCount = adjacencyListKeys.length @@ -119,7 +137,7 @@ export function getTransformedModelEdgesTargetSources< for (let i = 0; i < nodesCount; i++) { const targetAdjacencyListKey = adjacencyListKeys[i] - const targetNodeId = toNodeID(targetAdjacencyListKey) + const targetNodeId = toNodeID(targetAdjacencyListKey) const sources = lineageAdjacencyList[targetAdjacencyListKey] const sourcesCount = sources?.length || 0 @@ -134,7 +152,7 @@ export function getTransformedModelEdgesTargetSources< sourceAdjacencyListKey, targetAdjacencyListKey, ) - const sourceNodeId = toNodeID(sourceAdjacencyListKey) + const sourceNodeId = toNodeID(sourceAdjacencyListKey) edges.push(transformEdge('edge', edgeId, sourceNodeId, targetNodeId)) } @@ -206,18 +224,27 @@ export function calculateNodeDetailsHeight({ export function createEdge< TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, >( type: string, edgeId: TEdgeID, - sourceId: TNodeID, - targetId: TNodeID, - sourceHandleId?: TPortID, - targetHandleId?: TPortID, + sourceId: TSourceID, + targetId: TTargetID, + sourceHandleId?: TSourceHandleID, + targetHandleId?: TTargetHandleID, data?: TEdgeData, -): LineageEdge { +): LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID +> { return { id: edgeId, source: sourceId, diff --git a/web/common/src/components/Lineage/layout/dagreLayout.ts b/web/common/src/components/Lineage/layout/dagreLayout.ts index 83714a2220..554d427f03 100644 --- a/web/common/src/components/Lineage/layout/dagreLayout.ts +++ b/web/common/src/components/Lineage/layout/dagreLayout.ts @@ -13,14 +13,23 @@ import dagre from 'dagre' export function buildLayout< TNodeData extends LineageNodeData = LineageNodeData, TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, >({ edges, nodesMap, }: { - edges: LineageEdge[] + edges: LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >[] nodesMap: LineageNodesMap }) { const nodes = Object.values(nodesMap) diff --git a/web/common/src/components/Lineage/layout/help.ts b/web/common/src/components/Lineage/layout/help.ts index 91b3ebc4a3..d0dada83f5 100644 --- a/web/common/src/components/Lineage/layout/help.ts +++ b/web/common/src/components/Lineage/layout/help.ts @@ -24,14 +24,33 @@ export function getWorker(url: URL): Worker { export async function getLayoutedGraph< TNodeData extends LineageNodeData = LineageNodeData, TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, >( - edges: LineageEdge[], + edges: LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >[], nodesMap: LineageNodesMap, workerUrl: URL, -): Promise> { +): Promise< + LayoutedGraph< + TNodeData, + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + > +> { let timeoutId: NodeJS.Timeout | null = null return new Promise((resolve, reject) => { @@ -56,9 +75,11 @@ export async function getLayoutedGraph< worker.postMessage({ edges, nodesMap } as LayoutedGraph< TNodeData, TEdgeData, - TNodeID, TEdgeID, - TPortID + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID >) } catch (postError) { errorHandler(postError as ErrorEvent) @@ -66,7 +87,15 @@ export async function getLayoutedGraph< function handler( event: MessageEvent< - LayoutedGraph & { + LayoutedGraph< + TNodeData, + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + > & { error: ErrorEvent } >, diff --git a/web/common/src/components/Lineage/node/NodeHandle.tsx b/web/common/src/components/Lineage/node/NodeHandle.tsx index 4bfbfa6181..d50d90422a 100644 --- a/web/common/src/components/Lineage/node/NodeHandle.tsx +++ b/web/common/src/components/Lineage/node/NodeHandle.tsx @@ -3,8 +3,9 @@ import React from 'react' import { cn } from '@/utils' import { BaseHandle } from './base-handle' +import type { HandleId } from '../utils' -export const NodeHandle = React.memo(function NodeHandle({ +export function NodeHandle({ type, id, children, @@ -12,7 +13,7 @@ export const NodeHandle = React.memo(function NodeHandle({ ...props }: { type: 'target' | 'source' - id: string + id: THandleId children: React.ReactNode className?: string }) { @@ -29,4 +30,4 @@ export const NodeHandle = React.memo(function NodeHandle({ {children} ) -}) +} diff --git a/web/common/src/components/Lineage/node/NodeHandles.tsx b/web/common/src/components/Lineage/node/NodeHandles.tsx index 71bee716b4..453ff74317 100644 --- a/web/common/src/components/Lineage/node/NodeHandles.tsx +++ b/web/common/src/components/Lineage/node/NodeHandles.tsx @@ -3,8 +3,12 @@ import React from 'react' import { cn } from '@/utils' import { HorizontalContainer } from '@/components/HorizontalContainer/HorizontalContainer' import { NodeHandle } from './NodeHandle' +import type { HandleId } from '../utils' -export const NodeHandles = React.memo(function NodeHandles({ +export function NodeHandles< + TLeftHandleId extends string = HandleId, + TRightHandleId extends string = HandleId, +>({ leftIcon, rightIcon, leftId, @@ -13,8 +17,8 @@ export const NodeHandles = React.memo(function NodeHandles({ handleClassName, children, }: { - leftId?: string - rightId?: string + leftId?: TLeftHandleId + rightId?: TRightHandleId className?: string handleClassName?: string children: React.ReactNode @@ -27,7 +31,7 @@ export const NodeHandles = React.memo(function NodeHandles({ data-component="NodeHandles" > {leftId && ( - type="target" id={leftId} className={cn('left-0', handleClassName)} @@ -37,7 +41,7 @@ export const NodeHandles = React.memo(function NodeHandles({ )} {children} {rightId && ( - type="source" id={rightId} className={cn('right-0', handleClassName)} @@ -47,4 +51,4 @@ export const NodeHandles = React.memo(function NodeHandles({ )} ) -}) +} diff --git a/web/common/src/components/Lineage/node/NodePort.tsx b/web/common/src/components/Lineage/node/NodePort.tsx index b961d4e01a..7380716f02 100644 --- a/web/common/src/components/Lineage/node/NodePort.tsx +++ b/web/common/src/components/Lineage/node/NodePort.tsx @@ -2,12 +2,14 @@ import { useNodeConnections, useUpdateNodeInternals } from '@xyflow/react' import React from 'react' import { cn } from '@/utils' -import { type NodeId, type PortId } from '../utils' +import { type NodeId, type PortHandleId } from '../utils' import { NodeHandles } from './NodeHandles' -export const NodePort = React.memo(function NodePort< - TPortId extends string = PortId, +export function NodePort< + TPortId extends string = PortHandleId, TNodeID extends string = NodeId, + TLeftPortHandleId extends string = PortHandleId, + TRightPortHandleId extends string = PortHandleId, >({ id, nodeId, @@ -32,8 +34,16 @@ export const NodePort = React.memo(function NodePort< handleId: id, }) - const leftId = targets.length > 0 ? id : undefined - const rightId = sources.length > 0 ? id : undefined + const isLeftHandleId = (id: TPortId): id is TPortId & TLeftPortHandleId => { + return id && targets.length > 0 + } + + const isRightHandleId = (id: TPortId): id is TPortId & TRightPortHandleId => { + return id && sources.length > 0 + } + + const leftId = isLeftHandleId(id) ? id : undefined + const rightId = isRightHandleId(id) ? id : undefined React.useEffect(() => { if (leftId || rightId) { @@ -42,7 +52,7 @@ export const NodePort = React.memo(function NodePort< }, [updateNodeInternals, nodeId, leftId, rightId]) return ( - data-component="NodePort" leftIcon={ @@ -61,4 +71,4 @@ export const NodePort = React.memo(function NodePort< {children} ) -}) +} diff --git a/web/common/src/components/Lineage/node/base-handle.tsx b/web/common/src/components/Lineage/node/base-handle.tsx index 76d66bdeaf..e6b8f0c24b 100644 --- a/web/common/src/components/Lineage/node/base-handle.tsx +++ b/web/common/src/components/Lineage/node/base-handle.tsx @@ -16,7 +16,6 @@ export const BaseHandle: ForwardRefExoticComponent< 'fixed flex justify-center items-center border-none transition', className, )} - {...props} > {children} diff --git a/web/common/src/components/Lineage/stories/Lineage.stories.tsx b/web/common/src/components/Lineage/stories/Lineage.stories.tsx index 6e16bed61e..115be3c2c0 100644 --- a/web/common/src/components/Lineage/stories/Lineage.stories.tsx +++ b/web/common/src/components/Lineage/stories/Lineage.stories.tsx @@ -1,12 +1,125 @@ import type { LineageAdjacencyList, LineageDetails } from '../utils' import { ModelLineage } from './ModelLineage' -import type { ModelLineageNodeDetails, ModelName } from './ModelLineageContext' +import type { + BrandedLineageAdjacencyList, + BrandedLineageDetails, + ModelLineageNodeDetails, + ModelName, +} from './ModelLineageContext' export default { title: 'Components/Lineage', } +const adjacencyList = { + 'sqlmesh.sushi.raw_orders': ['sqlmesh.sushi.orders'], + 'sqlmesh.sushi.orders': [], +} as Record + +const lineageDetails = { + 'sqlmesh.sushi.raw_orders': { + name: 'sqlmesh.sushi.raw_orders', + display_name: 'sushi.raw_orders', + identifier: '123456789', + version: '123456789', + dialect: 'bigquery', + cron: '0 0 * * *', + owner: 'admin', + kind: 'INCREMENTAL_BY_TIME', + model_type: 'python', + tags: ['test', 'tag', 'another tag'], + columns: { + user_id: { + data_type: 'STRING', + description: 'node', + }, + event_id: { + data_type: 'STRING', + description: 'node', + }, + created_at: { + data_type: 'TIMESTAMP', + description: 'node', + }, + }, + }, + 'sqlmesh.sushi.orders': { + name: 'sqlmesh.sushi.orders', + display_name: 'sushi.orders', + identifier: '123456789', + version: '123456789', + dialect: 'bigquery', + cron: '0 0 * * *', + owner: 'admin', + kind: 'INCREMENTAL_BY_TIME', + model_type: 'sql', + tags: ['test', 'tag', 'another tag'], + columns: { + user_id: { + data_type: 'STRING', + description: 'node', + columnLineageData: { + 'sqlmesh.sushi.orders': { + user_id: { + source: 'sqlmesh.sushi.raw_orders', + expression: 'select user_id from sqlmesh.sushi.raw_orders', + models: { + 'sqlmesh.sushi.raw_orders': ['user_id'], + }, + }, + }, + }, + }, + event_id: { + data_type: 'STRING', + description: 'node', + columnLineageData: { + 'sqlmesh.sushi.orders': { + event_id: { + models: { + 'sqlmesh.sushi.raw_orders': ['event_id'], + }, + }, + }, + }, + }, + product_id: { + data_type: 'STRING', + description: 'node', + }, + customer_id: { + data_type: 'STRING', + description: 'node', + }, + updated_at: { + data_type: 'TIMESTAMP', + description: 'node', + }, + deleted_at: { + data_type: 'TIMESTAMP', + description: 'node', + }, + expired_at: { + data_type: 'TIMESTAMP', + description: 'node', + }, + start_at: { + data_type: 'TIMESTAMP', + description: 'node', + }, + end_at: { + data_type: 'TIMESTAMP', + description: 'node', + }, + created_ts: { + data_type: 'TIMESTAMP', + description: 'node', + }, + }, + }, +} as Record + export const LineageModel = () => { return (
{ `} - } - lineageDetails={ - { - 'sqlmesh.sushi.raw_orders': { - name: 'sqlmesh.sushi.raw_orders', - display_name: 'sushi.raw_orders', - identifier: '123456789', - version: '123456789', - dialect: 'bigquery', - cron: '0 0 * * *', - owner: 'admin', - kind: 'INCREMENTAL_BY_TIME', - model_type: 'python', - tags: ['test', 'tag', 'another tag'], - columns: { - user_id: { - data_type: 'STRING', - description: 'node', - }, - event_id: { - data_type: 'STRING', - description: 'node', - }, - created_at: { - data_type: 'TIMESTAMP', - description: 'node', - }, - }, - }, - 'sqlmesh.sushi.orders': { - name: 'sqlmesh.sushi.orders', - display_name: 'sushi.orders', - identifier: '123456789', - version: '123456789', - dialect: 'bigquery', - cron: '0 0 * * *', - owner: 'admin', - kind: 'INCREMENTAL_BY_TIME', - model_type: 'sql', - tags: ['test', 'tag', 'another tag'], - columns: { - user_id: { - data_type: 'STRING', - description: 'node', - columnLineageData: { - 'sqlmesh.sushi.orders': { - user_id: { - source: 'sqlmesh.sushi.raw_orders', - expression: - 'select user_id from sqlmesh.sushi.raw_orders', - models: { - 'sqlmesh.sushi.raw_orders': ['user_id'], - }, - }, - }, - }, - }, - event_id: { - data_type: 'STRING', - description: 'node', - columnLineageData: { - 'sqlmesh.sushi.orders': { - event_id: { - models: { - 'sqlmesh.sushi.raw_orders': ['event_id'], - }, - }, - }, - }, - }, - product_id: { - data_type: 'STRING', - description: 'node', - }, - customer_id: { - data_type: 'STRING', - description: 'node', - }, - updated_at: { - data_type: 'TIMESTAMP', - description: 'node', - }, - deleted_at: { - data_type: 'TIMESTAMP', - description: 'node', - }, - expired_at: { - data_type: 'TIMESTAMP', - description: 'node', - }, - start_at: { - data_type: 'TIMESTAMP', - description: 'node', - }, - end_at: { - data_type: 'TIMESTAMP', - description: 'node', - }, - created_ts: { - data_type: 'TIMESTAMP', - description: 'node', - }, - }, - }, - } as LineageDetails - } + adjacencyList={adjacencyList as BrandedLineageAdjacencyList} + lineageDetails={lineageDetails as BrandedLineageDetails} className="rounded-2xl" />
diff --git a/web/common/src/components/Lineage/stories/ModelLineage.tsx b/web/common/src/components/Lineage/stories/ModelLineage.tsx index 46d19f9758..3df85ea1a5 100644 --- a/web/common/src/components/Lineage/stories/ModelLineage.tsx +++ b/web/common/src/components/Lineage/stories/ModelLineage.tsx @@ -1,8 +1,6 @@ -import { debounce } from 'lodash' import { Focus, LockOpen, Rows2, Rows3, Lock } from 'lucide-react' import React from 'react' -import { type ColumnLevelLineageAdjacencyList } from '../LineageColumnLevel/ColumnLevelLineageContext' import { MAX_COLUMNS_TO_DISPLAY, calculateColumnsHeight, @@ -11,16 +9,8 @@ import { getEdgesFromColumnLineage, } from '../LineageColumnLevel/help' import { useColumnLevelLineage } from '../LineageColumnLevel/useColumnLevelLineage' -import { LineageControlButton } from '../LineageControlButton' -import { LineageControlIcon } from '../LineageControlIcon' import { LineageLayout } from '../LineageLayout' import { FactoryEdgeWithGradient } from '../edge/FactoryEdgeWithGradient' -import { - toNodeID, - toPortID, - type LineageAdjacencyList, - type LineageDetails, -} from '../utils' import { calculateNodeBaseHeight, calculateNodeDetailsHeight, @@ -33,6 +23,8 @@ import { import { type LineageEdge, type LineageNodesMap, + toNodeID, + toPortID, ZOOM_THRESHOLD, } from '../utils' import { @@ -47,11 +39,23 @@ import { type ModelColumnID, type ModelEdgeId, type NodeType, + type BrandedLineageAdjacencyList, + type BrandedLineageDetails, + type BrandedColumnLevelLineageAdjacencyList, + type ModelColumn, + type ModelDisplayName, + type LeftHandleId, + type RightHandleId, + type LeftPortHandleId, + type RightPortHandleId, } from './ModelLineageContext' import { ModelNode } from './ModelNode' import { getNodeTypeColorVar } from './help' import { EdgeWithGradient } from '../edge/EdgeWithGradient' import { cleanupLayoutWorker, getLayoutedGraph } from '../layout/help' +import { LineageControlButton } from '../LineageControlButton' +import { LineageControlIcon } from '../LineageControlIcon' +import type { BrandedRecord } from '@/types' const nodeTypes = { node: ModelNode, @@ -67,8 +71,8 @@ export const ModelLineage = ({ lineageDetails, className, }: { - adjacencyList: LineageAdjacencyList - lineageDetails: LineageDetails + adjacencyList: BrandedLineageAdjacencyList + lineageDetails: BrandedLineageDetails selectedModelName?: ModelName className?: string }) => { @@ -76,7 +80,14 @@ export const ModelLineage = ({ const [isBuildingLayout, setIsBuildingLayout] = React.useState(false) const [nodesDraggable, setNodesDraggable] = React.useState(false) const [edges, setEdges] = React.useState< - LineageEdge[] + LineageEdge< + EdgeData, + ModelEdgeId, + LeftHandleId, + RightHandleId, + LeftPortHandleId, + RightPortHandleId + >[] >([]) const [nodesMap, setNodesMap] = React.useState< LineageNodesMap @@ -94,7 +105,7 @@ export const ModelLineage = ({ const [showColumns, setShowColumns] = React.useState(false) const [columnLevelLineage, setColumnLevelLineage] = React.useState< - Map> + Map >(new Map()) const [fetchingColumns, setFetchingColumns] = React.useState< Set @@ -104,9 +115,12 @@ export const ModelLineage = ({ adjacencyListColumnLevel, selectedColumns, adjacencyListKeysColumnLevel, - } = useColumnLevelLineage( - columnLevelLineage, - ) + } = useColumnLevelLineage< + ModelName, + ColumnName, + ModelColumnID, + BrandedColumnLevelLineageAdjacencyList + >(columnLevelLineage) const adjacencyListKeys = React.useMemo(() => { let keys: ModelName[] = [] @@ -124,18 +138,18 @@ export const ModelLineage = ({ (nodeId: ModelNodeId, detail: ModelLineageNodeDetails) => { const columns = detail.columns - const node = createNode('node', nodeId, { - name: detail.name, + const node = createNode('node', nodeId, { + name: detail.name as ModelName, + displayName: detail.display_name as ModelDisplayName, identifier: detail.identifier, model_type: detail.model_type as NodeType, kind: detail.kind!, cron: detail.cron, - displayName: detail.display_name, owner: detail.owner!, dialect: detail.dialect, version: detail.version, tags: detail.tags || [], - columns, + columns: columns as BrandedRecord, }) const selectedColumnsCount = new Set( Object.keys(columns ?? {}).map(k => toPortID(detail.name, k)), @@ -184,10 +198,10 @@ export const ModelLineage = ({ ( edgeType: string, edgeId: ModelEdgeId, - sourceId: ModelNodeId, - targetId: ModelNodeId, - sourceHandleId?: ModelColumnID, - targetHandleId?: ModelColumnID, + sourceId: LeftHandleId, + targetId: RightHandleId, + sourceHandleId?: LeftPortHandleId, + targetHandleId?: RightPortHandleId, ) => { const sourceNode = transformedNodesMap[sourceId] const targetNode = transformedNodesMap[targetId] @@ -217,7 +231,14 @@ export const ModelLineage = ({ data.strokeWidth = 2 } - return createEdge( + return createEdge< + EdgeData, + ModelEdgeId, + LeftHandleId, + RightHandleId, + LeftPortHandleId, + RightPortHandleId + >( edgeType, edgeId, sourceId, @@ -237,8 +258,11 @@ export const ModelLineage = ({ ColumnName, EdgeData, ModelEdgeId, - ModelNodeId, - ModelColumnID + LeftHandleId, + RightHandleId, + LeftPortHandleId, + RightPortHandleId, + BrandedColumnLevelLineageAdjacencyList >({ columnLineage: adjacencyListColumnLevel, transformEdge, @@ -252,38 +276,45 @@ export const ModelLineage = ({ : getTransformedModelEdgesSourceTargets< ModelName, EdgeData, - ModelNodeId, ModelEdgeId, - ModelColumnID + LeftHandleId, + RightHandleId, + LeftPortHandleId, + RightPortHandleId >(adjacencyListKeys, adjacencyList, transformEdge) }, [adjacencyListKeys, adjacencyList, transformEdge, edgesColumnLevel]) - const calculateLayout = React.useMemo(() => { - return debounce( - ( - eds: LineageEdge[], - nds: LineageNodesMap, - ) => - getLayoutedGraph( - eds, - nds, - new URL('./dagreLayout.worker.ts', import.meta.url), - ) - .then(({ edges, nodesMap }) => { - setEdges(edges) - setNodesMap(nodesMap) - }) - .catch(error => { - console.error('Layout processing failed:', error) - setEdges([]) - setNodesMap({}) - }) - .finally(() => { - setIsBuildingLayout(false) - }), - 200, - ) - }, []) + const calculateLayout = React.useCallback( + ( + eds: LineageEdge< + EdgeData, + ModelEdgeId, + LeftHandleId, + RightHandleId, + LeftPortHandleId, + RightPortHandleId + >[], + nds: LineageNodesMap, + ) => + getLayoutedGraph( + eds, + nds, + new URL('./dagreLayout.worker.ts', import.meta.url), + ) + .then(({ edges, nodesMap }) => { + setEdges(edges) + setNodesMap(nodesMap) + }) + .catch(error => { + console.error('Layout processing failed:', error) + setEdges([]) + setNodesMap({}) + }) + .finally(() => { + setIsBuildingLayout(false) + }), + [setEdges, setNodesMap, setIsBuildingLayout], + ) const nodes = React.useMemo(() => { return Object.values(nodesMap) @@ -291,7 +322,7 @@ export const ModelLineage = ({ const currentNode = React.useMemo(() => { return selectedModelName - ? nodesMap[toNodeID(selectedModelName)] + ? nodesMap[toNodeID(selectedModelName)] : null }, [selectedModelName, nodesMap]) @@ -315,30 +346,13 @@ export const ModelLineage = ({ selectedNodes, ) const onlySelectedEdges = transformedEdges.filter(edge => - selectedEdges.has(edge.id as ModelEdgeId), + selectedEdges.has(edge.id), ) calculateLayout(onlySelectedEdges, onlySelectedNodesMap) } else { calculateLayout(transformedEdges, transformedNodesMap) } - }, [ - calculateLayout, - showOnlySelectedNodes, - transformedEdges, - transformedNodesMap, - ]) - - React.useEffect(() => { - const currentNodeId = selectedModelName - ? toNodeID(selectedModelName) - : undefined - - if (currentNodeId && currentNodeId in nodesMap) { - setSelectedNodeId(currentNodeId) - } else { - handleReset() - } - }, [handleReset, selectedModelName, nodesMap]) + }, [showOnlySelectedNodes, transformedEdges, transformedNodesMap]) // Cleanup worker on unmount React.useEffect(() => () => cleanupLayoutWorker(), []) @@ -381,7 +395,10 @@ export const ModelLineage = ({ EdgeData, ModelNodeId, ModelEdgeId, - ModelColumnID + LeftHandleId, + RightHandleId, + LeftPortHandleId, + RightPortHandleId > isBuildingLayout={isBuildingLayout} useLineage={useModelLineage} diff --git a/web/common/src/components/Lineage/stories/ModelLineageContext.ts b/web/common/src/components/Lineage/stories/ModelLineageContext.ts index 98d2131766..39e23ceb6d 100644 --- a/web/common/src/components/Lineage/stories/ModelLineageContext.ts +++ b/web/common/src/components/Lineage/stories/ModelLineageContext.ts @@ -1,4 +1,4 @@ -import type { Branded } from '@/types' +import type { Branded, BrandedRecord } from '@/types' import { type ColumnLevelLineageAdjacencyList, type ColumnLevelLineageContextValue, @@ -10,22 +10,54 @@ import { createLineageContext, getInitial as getLineageContextInitial, } from '../LineageContext' -import { type PathType } from '../utils' +import { + type LineageAdjacencyList, + type LineageDetails, + type PathType, +} from '../utils' export type ModelName = Branded +export type ModelDisplayName = Branded + export type ColumnName = Branded + export type ModelColumnID = Branded -export type ModelNodeId = Branded export type ModelEdgeId = Branded + +export type LeftHandleId = Branded +export type RightHandleId = Branded + +export type LeftPortHandleId = Branded +export type RightPortHandleId = Branded + +export type ModelNodeId = LeftHandleId | RightHandleId + +export type BrandedColumnLevelLineageAdjacencyList = + ColumnLevelLineageAdjacencyList & { + readonly __adjacencyListKeyBrand: ModelName + readonly __adjacencyListColumnKeyBrand: ColumnName + } + +export type BrandedLineageAdjacencyList = LineageAdjacencyList & { + readonly __adjacencyListKeyBrand: ModelName +} + +export type BrandedLineageDetails = LineageDetails< + ModelName, + ModelLineageNodeDetails +> & { + readonly __lineageDetailsKeyBrand: ModelName +} + export type ModelColumn = Column & { id: ModelColumnID name: ColumnName - columnLineageData?: ColumnLevelLineageAdjacencyList + columnLineageData?: BrandedColumnLevelLineageAdjacencyList } export type NodeType = 'sql' | 'python' export type ModelLineageNodeDetails = { - name: ModelName + name: string display_name: string identifier: string version: string @@ -35,12 +67,12 @@ export type ModelLineageNodeDetails = { kind?: string model_type?: string tags?: string[] - columns?: Record + columns?: BrandedRecord } export type NodeData = { name: ModelName - displayName: string + displayName: ModelDisplayName model_type: NodeType identifier: string version: string @@ -48,8 +80,8 @@ export type NodeData = { cron: string owner: string dialect: string - columns?: Record tags: string[] + columns?: BrandedRecord } export type EdgeData = { @@ -62,14 +94,18 @@ export type EdgeData = { export type ModelLineageContextValue = ColumnLevelLineageContextValue< ModelName, ColumnName, - ModelColumnID + ModelColumnID, + BrandedColumnLevelLineageAdjacencyList > & LineageContextValue< NodeData, EdgeData, ModelNodeId, ModelEdgeId, - ModelColumnID + LeftHandleId, + RightHandleId, + LeftPortHandleId, + RightPortHandleId > export const initial = { @@ -77,7 +113,8 @@ export const initial = { ...getColumnLevelLineageContextInitial< ModelName, ColumnName, - ModelColumnID + ModelColumnID, + BrandedColumnLevelLineageAdjacencyList >(), } @@ -86,7 +123,10 @@ export const { Provider, useLineage } = createLineageContext< EdgeData, ModelNodeId, ModelEdgeId, - ModelColumnID, + LeftHandleId, + RightHandleId, + LeftPortHandleId, + RightPortHandleId, ModelLineageContextValue >(initial) diff --git a/web/common/src/components/Lineage/stories/ModelNodeColumn.tsx b/web/common/src/components/Lineage/stories/ModelNodeColumn.tsx index 35d4a0e592..dbb3f92dad 100644 --- a/web/common/src/components/Lineage/stories/ModelNodeColumn.tsx +++ b/web/common/src/components/Lineage/stories/ModelNodeColumn.tsx @@ -1,6 +1,5 @@ import React from 'react' -import { type ColumnLevelLineageAdjacencyList } from '../LineageColumnLevel/ColumnLevelLineageContext' import { FactoryColumn } from '../LineageColumnLevel/FactoryColumn' import { @@ -9,13 +8,19 @@ import { type ModelName, type ModelNodeId, type ColumnName, + type BrandedColumnLevelLineageAdjacencyList, + type LeftPortHandleId, + type RightPortHandleId, } from './ModelLineageContext' const ModelColumn = FactoryColumn< ModelName, ColumnName, ModelNodeId, - ModelColumnID + ModelColumnID, + LeftPortHandleId, + RightPortHandleId, + BrandedColumnLevelLineageAdjacencyList >(useModelLineage) export const ModelNodeColumn = React.memo(function ModelNodeColumn({ @@ -35,7 +40,7 @@ export const ModelNodeColumn = React.memo(function ModelNodeColumn({ type: string description?: string | null className?: string - columnLineageData?: ColumnLevelLineageAdjacencyList + columnLineageData?: BrandedColumnLevelLineageAdjacencyList }) { const { selectedColumns, setColumnLevelLineage } = useModelLineage() diff --git a/web/common/src/components/Lineage/utils.ts b/web/common/src/components/Lineage/utils.ts index 01a277f17a..4e2d55a5a0 100644 --- a/web/common/src/components/Lineage/utils.ts +++ b/web/common/src/components/Lineage/utils.ts @@ -4,6 +4,8 @@ import { type Edge, type Node } from '@xyflow/react' export type NodeId = Branded export type EdgeId = Branded export type PortId = Branded +export type HandleId = Branded +export type PortHandleId = Branded export type LineageNodeData = Record export type LineageEdgeData = Record @@ -29,25 +31,36 @@ export interface LineageNode< export interface LineageEdge< TEdgeData extends LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, > extends Edge { id: TEdgeID - source: TNodeID - target: TNodeID - sourceHandle?: TPortID - targetHandle?: TPortID + source: TSourceID + target: TTargetID + sourceHandle?: TSourceHandleID + targetHandle?: TTargetHandleID } export type LayoutedGraph< TNodeData extends LineageNodeData = LineageNodeData, TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, > = { - edges: LineageEdge[] + edges: LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID + >[] nodesMap: LineageNodesMap } @@ -60,17 +73,26 @@ export type TransformNodeFn< export type TransformEdgeFn< TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, TEdgeID extends string = EdgeId, - TPortID extends string = PortId, + TSourceID extends string = NodeId, + TTargetID extends string = NodeId, + TSourceHandleID extends string = PortId, + TTargetHandleID extends string = PortId, > = ( edgeType: string, edgeId: TEdgeID, - sourceId: TNodeID, - targetId: TNodeID, - sourceColumnId?: TPortID, - targetColumnId?: TPortID, -) => LineageEdge + sourceId: TSourceID, + targetId: TTargetID, + sourceHandleId?: TSourceHandleID, + targetHandleId?: TTargetHandleID, +) => LineageEdge< + TEdgeData, + TEdgeID, + TSourceID, + TTargetID, + TSourceHandleID, + TTargetHandleID +> export const DEFAULT_NODE_HEIGHT = 32 export const DEFAULT_NODE_WIDTH = 300 diff --git a/web/common/src/index.ts b/web/common/src/index.ts index 0748a6c78e..c3c65a8e77 100644 --- a/web/common/src/index.ts +++ b/web/common/src/index.ts @@ -66,6 +66,8 @@ export { cn, truncate } from '@/utils' export type { Brand, Branded, + BrandedString, + BrandedRecord, Size, HeadlineLevel, Side, diff --git a/web/common/src/types.ts b/web/common/src/types.ts index 3de26b205d..e8bdf3e9de 100644 --- a/web/common/src/types.ts +++ b/web/common/src/types.ts @@ -1,8 +1,46 @@ export declare const __brand: unique symbol - export type Brand = { [__brand]: B } + +/** + * Branded is a type that adds a brand to a type. It is a type that is used to + * ensure that the type is unique and that it is not possible to mix up types + * with the same brand. + * + * @example + * + * type UserId = Branded + * type UserName = Branded + * + * const userId = '123' as UserId + * const userName = 'John Doe' as UserName + * + * userId == userName -> compile error + */ export type Branded = T & Brand +/** + * Constraint that only accepts branded string types + */ +export type BrandedString = string & Brand + +/** + * BrandedRecord is a type that creates a branded Record type with strict key checking. + * This ensures that Record is NOT assignable to Record + * + * @example + * type ModelFQN = Branded + * type ModelName = Branded + * + * type FQNMap = BrandedRecord + * type NameMap = BrandedRecord + * + * const fqnMap: FQNMap = {} + * const nameMap: NameMap = fqnMap // TypeScript error! + */ +export type BrandedRecord = Record & { + readonly __recordKeyBrand: K +} + export type Callback = (data?: T) => void export type Size = '2xs' | 'xs' | 's' | 'm' | 'l' | 'xl' | '2xl' From 3853a85606c030dec4cf45e8bfde1c033753385c Mon Sep 17 00:00:00 2001 From: Max Mykal Date: Tue, 7 Oct 2025 18:04:34 -0700 Subject: [PATCH 2/2] celan up --- .../src/components/Lineage/stories/ModelLineageContext.ts | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/web/common/src/components/Lineage/stories/ModelLineageContext.ts b/web/common/src/components/Lineage/stories/ModelLineageContext.ts index 39e23ceb6d..745d9c2636 100644 --- a/web/common/src/components/Lineage/stories/ModelLineageContext.ts +++ b/web/common/src/components/Lineage/stories/ModelLineageContext.ts @@ -18,20 +18,15 @@ import { export type ModelName = Branded export type ModelDisplayName = Branded - export type ColumnName = Branded - export type ModelColumnID = Branded export type ModelEdgeId = Branded - export type LeftHandleId = Branded export type RightHandleId = Branded - +export type ModelNodeId = LeftHandleId | RightHandleId export type LeftPortHandleId = Branded export type RightPortHandleId = Branded -export type ModelNodeId = LeftHandleId | RightHandleId - export type BrandedColumnLevelLineageAdjacencyList = ColumnLevelLineageAdjacencyList & { readonly __adjacencyListKeyBrand: ModelName