sqlmesh/web/common/src/components/Lineage/LineageContext.ts at 2d485cc43599cfb96f7d02c668f94f791824568d · SQLMesh/sqlmesh · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import React from 'react'

import {
  type EdgeId,
  type LineageEdge,
  type LineageEdgeData,
  type LineageNode,
  type LineageNodeData,
  type LineageNodesMap,
  type NodeId,
  type PortId,
  ZOOM_THRESHOLD,
} from './utils'

export interface LineageContextValue<
  TNodeData extends LineageNodeData = LineageNodeData,
  TEdgeData extends LineageEdgeData = LineageEdgeData,
  TNodeID extends string = NodeId,
  TEdgeID extends string = EdgeId,
  TPortID extends string = PortId,
> {
  // Node selection
  showOnlySelectedNodes: boolean
  setShowOnlySelectedNodes: React.Dispatch<React.SetStateAction<boolean>>
  selectedNodes: Set<TNodeID>
  setSelectedNodes: React.Dispatch<React.SetStateAction<Set<TNodeID>>>
  selectedEdges: Set<TEdgeID>
  setSelectedEdges: React.Dispatch<React.SetStateAction<Set<TEdgeID>>>
  selectedNodeId: TNodeID | null
  setSelectedNodeId: React.Dispatch<React.SetStateAction<TNodeID | null>>

  // Layout
  zoom: number
  setZoom: React.Dispatch<React.SetStateAction<number>>

  // Nodes and Edges
  edges: LineageEdge<TEdgeData, TNodeID, TEdgeID, TPortID>[]
  setEdges: React.Dispatch<
    React.SetStateAction<LineageEdge<TEdgeData, TNodeID, TEdgeID, TPortID>[]>
  >
  nodes: LineageNode<TNodeData, TNodeID>[]
  nodesMap: LineageNodesMap<TNodeData, TNodeID>
  setNodesMap: React.Dispatch<React.SetStateAction<LineageNodesMap<TNodeData>>>
  currentNode: LineageNode<TNodeData, TNodeID> | null
}

export function getInitial<
  TNodeID extends string = NodeId,
  TEdgeID extends string = EdgeId,
>() {
  return {
    showOnlySelectedNodes: false,
    setShowOnlySelectedNodes: () => {},
    selectedNodes: new Set<TNodeID>(),
    setSelectedNodes: () => {},
    selectedEdges: new Set<TEdgeID>(),
    setSelectedEdges: () => {},
    selectedNodeId: null,
    setSelectedNodeId: () => {},
    zoom: ZOOM_THRESHOLD,
    setZoom: () => {},
    edges: [],
    setEdges: () => {},
    nodes: [],
    nodesMap: {},
    setNodesMap: () => {},
    currentNode: null,
  }
}

export type LineageContextHook<
  TNodeData extends LineageNodeData = LineageNodeData,
  TEdgeData extends LineageEdgeData = LineageEdgeData,
  TNodeID extends string = NodeId,
  TEdgeID extends string = EdgeId,
  TPortID extends string = PortId,
> = () => LineageContextValue<TNodeData, TEdgeData, TNodeID, TEdgeID, TPortID>

export function createLineageContext<
  TNodeData extends LineageNodeData = LineageNodeData,
  TEdgeData extends LineageEdgeData = LineageEdgeData,
  TNodeID extends string = NodeId,
  TEdgeID extends string = EdgeId,
  TPortID extends string = PortId,
  TLineageContextValue extends LineageContextValue<
    TNodeData,
    TEdgeData,
    TNodeID,
    TEdgeID,
    TPortID
  > = LineageContextValue<TNodeData, TEdgeData, TNodeID, TEdgeID, TPortID>,
>(initial: TLineageContextValue) {
  const LineageContext = React.createContext(initial)

  return {
    Provider: LineageContext.Provider,
    useLineage: () => React.useContext(LineageContext),
  }
}