chore(vscode): further typing nodes

[ci skip]

- last bit of typing the logic

Author: benfdking

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

@@ -30,24 +30,24 @@ import { useLineageFlow } from './context'
 import { useApiColumnLineage } from '@/api/index'
 import SourceList from '@/components/sourceList/SourceList'
 import type { Lineage } from '@/domain/lineage'
-import type { Column, ColumnName } from '@/domain/column'
-import type { ModelEncodedFQN } from '@/domain/models'
+import type { Column, ColumnEncodedName, ColumnName } from '@/domain/column'
+import { encode, type ModelEncodedFQN } from '@/domain/models'
 
 export function ModelColumns({
   nodeId,
   columns,
   disabled,
   className,
-  limit = 5,
+  limit,
   withHandles = false,
   withDescription = true,
   maxHeight = '50vh',
 }: {
   nodeId: ModelEncodedFQN
   columns: Column[]
+  limit: number
   disabled?: boolean
   className?: string
-  limit?: number
   withHandles?: boolean
   withDescription?: boolean
   maxHeight?: string
@@ -292,10 +292,12 @@ export function ModelColumns({
                 withHandles={withHandles}
                 withDescription={withDescription}
                 isEmpty={
-                  isNotNil(getColumnFromLineage(lineage, nodeId, item.name)) &&
+                  isNotNil(
+                    getColumnFromLineage(lineage, nodeId, encode(item.name)),
+                  ) &&
                   Object.keys(
-                    getColumnFromLineage(lineage, nodeId, item.name)?.models ??
-                      {},
+                    getColumnFromLineage(lineage, nodeId, encode(item.name))
+                      ?.models ?? {},
                   ).length === 0
                 }
                 className="border-t border-neutral-10 first:border-0 rounded-md"
@@ -585,9 +587,9 @@ function ColumnStatus({
 }
 
 function getColumnFromLineage(
-  lineage: Record<string, Lineage>,
-  nodeId: string,
-  columnName: string,
+  lineage: Record<ModelEncodedFQN, Lineage>,
+  nodeId: ModelEncodedFQN,
+  columnName: ColumnEncodedName,
 ): LineageColumn | undefined {
-  return lineage?.[nodeId]?.columns?.[columnName as ColumnName]
+  return lineage?.[nodeId]?.columns?.[columnName]
 }

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

@@ -3,7 +3,7 @@ import clsx from 'clsx'
 import { useMemo, useCallback, useState } from 'react'
 import { ModelType, type Model } from '@/api/client'
 import { useLineageFlow } from './context'
-import { type GraphNodeData } from './help'
+import { MAX_VISIBLE_COLUMNS, type GraphNodeData } from './help'
 import { Position, type NodeProps } from 'reactflow'
 import { ModelNodeHeaderHandles } from './ModelNodeHeaderHandles'
 import { ModelColumns } from './ModelColumns'
@@ -37,7 +37,6 @@ export default function ModelNode({
     withColumns: false,
   }
   const {
-    // connections,
     models,
     handleClickModel,
     lineage,
@@ -229,6 +228,7 @@ export default function ModelNode({
           className="nowheel rounded-b-lg bg-theme-lighter text-xs"
           nodeId={id}
           columns={columns}
+          limit={MAX_VISIBLE_COLUMNS}
           disabled={shouldDisableColumns}
           withHandles={true}
           withDescription={false}

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

@@ -7,6 +7,7 @@ import { EnumSide, toID } from './types'
 import { ArrowRightCircleIcon } from '@heroicons/react/24/solid'
 import clsx from 'clsx'
 import { type LineageNodeModelType } from './ModelNode'
+import type { ModelEncodedFQN } from '@/domain/models'
 
 export function ModelNodeHeaderHandles({
   id,
@@ -21,7 +22,7 @@ export function ModelNodeHeaderHandles({
   handleClick,
   handleSelect,
 }: {
-  id: string
+  id: ModelEncodedFQN
   label: string
   type?: LineageNodeModelType
   hasLeft?: boolean

vscode/react/src/components/graph/help.ts

@@ -6,36 +6,40 @@ import {
   isNotNil,
   isObjectEmpty,
 } from '@/utils/index'
-import { type LineageColumn } from '@/api/client'
 import { Position, type Edge, type Node, type XYPosition } from 'reactflow'
 import { type ActiveEdges, type Connections } from './context'
 import { EnumSide, toID, toKeys } from './types'
 import {
   EnumLineageNodeModelType,
   type LineageNodeModelType,
 } from './ModelNode'
-import type { Lineage } from '@/domain/lineage'
+import type { Lineage, LineageColumn } from '@/domain/lineage'
 import type { ConnectedNode } from '@/workers/lineage'
-import { encode, type ModelEncodedFQN, type ModelURI } from '@/domain/models'
+import {
+  encode,
+  type ModelEncodedFQN,
+  type ModelName,
+  type ModelURI,
+} from '@/domain/models'
 import type { Column, ColumnName } from '@/domain/column'
 import type { ModelSQLMeshModel } from '@/domain/sqlmesh-model'
 
 /**
  * Space between nodes.
  */
-const NODE_BALANCE_SPACE = 64
+export const NODE_BALANCE_SPACE = 64
 /**
  * Height of a column line.
  */
-const COLUMN_LINE_HEIGHT = 24
+export const COLUMN_LINE_HEIGHT = 24
 /**
  * Assumed width of a character.
  */
-const CHAR_WIDTH = 8
+export const CHAR_WIDTH = 8
 /**
  * Maximum number of columns that can be visible in a node.
  */
-const MAX_VISIBLE_COLUMNS = 5
+export const MAX_VISIBLE_COLUMNS = 5
 
 export interface GraphNodeData {
   label: string
@@ -327,10 +331,10 @@ function createGraphEdge<Data>(
 
 export function mergeLineageWithColumns(
   currentLineage: Record<string, Lineage> = {},
-  newLineage: Record<string, Record<string, LineageColumn>> = {},
+  newLineage: Record<ModelName, Record<ColumnName, LineageColumn>> = {},
 ): Record<string, Lineage> {
-  for (const targetModelName in newLineage) {
-    const targetModelNameEncoded = encodeURI(targetModelName)
+  for (const targetModelName of toKeys(newLineage)) {
+    const targetModelNameEncoded = encode(targetModelName)
 
     if (isNil(currentLineage[targetModelNameEncoded])) {
       currentLineage[targetModelNameEncoded] = { columns: {}, models: [] }
@@ -339,16 +343,16 @@ export function mergeLineageWithColumns(
     const currentLineageModel = currentLineage[targetModelNameEncoded]!
     const newLineageModel = newLineage[targetModelName]!
 
-    for (const targetColumnName in newLineageModel) {
-      const targetColumnNameEncoded = encodeURI(targetColumnName)
+    for (const targetColumnName of toKeys(newLineageModel)) {
+      const targetColumnNameEncoded = encode(targetColumnName)
       const newLineageModelColumn = newLineageModel[targetColumnName]!
 
       if (isNil(currentLineageModel.columns)) {
         currentLineageModel.columns = {}
       }
 
       // New Column Lineage delivers fresh data, so we can just assign it
-      currentLineageModel.columns[targetColumnNameEncoded as ColumnName] = {
+      currentLineageModel.columns[targetColumnNameEncoded] = {
         expression: newLineageModelColumn.expression ?? undefined,
         source: newLineageModelColumn.source ?? undefined,
         models: {},
@@ -362,7 +366,7 @@ export function mergeLineageWithColumns(
       const currentLineageModelColumnModels = currentLineageModelColumn.models
 
       for (const sourceColumnName in newLineageModelColumn.models) {
-        const sourceColumnNameEncoded = encodeURI(sourceColumnName)
+        const sourceColumnNameEncoded = encode(sourceColumnName)
         const currentLineageModelColumnModel =
           currentLineageModelColumnModels[
             sourceColumnNameEncoded as ModelEncodedFQN
@@ -400,11 +404,11 @@ export function mergeConnections(
 
   // We are getting lineage in format of target -> source
   for (const targetModelName in lineage) {
-    const targetModelNameEncoded = encodeURI(targetModelName)
+    const targetModelNameEncoded = encode(targetModelName)
     const model = lineage[targetModelName]!
 
     for (const targetColumnName in model) {
-      const targetColumnNameEncoded = encodeURI(targetColumnName)
+      const targetColumnNameEncoded = encode(targetColumnName)
       const column = model[targetColumnName]
 
       // We don't have any connectins so we skip
@@ -426,9 +430,9 @@ export function mergeConnections(
       }
 
       Object.entries(column.models).forEach(([sourceModelName, columns]) => {
-        const sourceModelNameEncoded = encodeURI(sourceModelName)
+        const sourceModelNameEncoded = encode(sourceModelName)
         columns.forEach(sourceColumnName => {
-          const sourceColumnNameEncoded = encodeURI(sourceColumnName)
+          const sourceColumnNameEncoded = encode(sourceColumnName)
           // It is a source (right handler)
           // but it can also be a target (left handler) for other connections
           const modelColumnIdSource = toID(

vscode/react/src/components/graph/types.ts

@@ -1,4 +1,4 @@
-import type { ColumnName } from '@/domain/column'
+import type { ColumnEncodedName } from '@/domain/column'
 import type { ModelEncodedFQN } from '@/domain/models'
 
 export const EnumSide = {
@@ -21,7 +21,7 @@ export type EdgeId = string
 export function toID(
   leftOrRight: 'left' | 'right',
   modelName: ModelEncodedFQN,
-  columnName: ColumnName,
+  columnName: ColumnEncodedName,
 ): NodeId
 export function toID(source: NodeId, target: NodeId): NodeId
 export function toID(

vscode/react/src/domain/column.ts

@@ -3,6 +3,8 @@ import { type Branded } from '@bus/brand'
 
 export type ColumnName = Branded<string, 'ColumnName'>
 
+export type ColumnEncodedName = Branded<string, 'ColumnEncodedName'>
+
 export type Column = {
   name: ColumnName
   type: string

vscode/react/src/domain/lineage.ts

@@ -3,11 +3,11 @@ import {
   type ModelLineageApiLineageModelNameGet200,
 } from '@/api/client'
 import type { ModelEncodedFQN, ModelFQN } from '@/domain/models'
-import type { ColumnName } from './column'
+import type { ColumnEncodedName, ColumnName } from './column'
 
 export interface Lineage {
   models: ModelEncodedFQN[]
-  columns?: Record<ColumnName, LineageColumn>
+  columns?: Record<ColumnEncodedName, LineageColumn>
 }
 
 export interface LineageColumn {

vscode/react/src/domain/models.ts

@@ -1,4 +1,5 @@
 import type { Branded } from '@bus/brand'
+import type { ColumnEncodedName, ColumnName } from './column'
 
 /**
  * ModelName is a type that represents the name of a model.
@@ -30,13 +31,15 @@ export type ModelURI = Branded<string, 'ModelURI'>
  */
 export type ModelEncodedURI = Branded<string, 'ModelEncodedURI'>
 
-export function encode(fqn: ModelName): ModelEncodedName
-export function encode(fqn: ModelURI): ModelEncodedURI
+export function encode(columnName: ColumnName): ColumnEncodedName
+export function encode(name: ModelName): ModelEncodedName
+export function encode(uri: ModelURI): ModelEncodedURI
 export function encode(fqn: ModelFQN): ModelEncodedFQN
 export function encode(s: string): string {
   return encodeURI(s)
 }
 
+export function decode(columnName: ColumnEncodedName): ColumnName
 export function decode(fqn: ModelEncodedName): ModelName
 export function decode(fqn: ModelEncodedURI): ModelURI
 export function decode(fqn: ModelEncodedFQN): ModelFQN