chore: add branded model string types

Author: benfdking

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

@@ -46,6 +46,7 @@ import { Popover, Transition } from '@headlessui/react'
 import { useApiColumnLineage } from '@/api/index'
 import SourceList from '@/components/sourceList/SourceList'
 import type { Lineage } from '@/domain/lineage'
+import type { ModelName } from '@/types/models'
 
 export default function ModelColumns({
   nodeId,
@@ -732,5 +733,5 @@ function getColumnFromLineage(
   nodeId: string,
   columnName: string,
 ): LineageColumn | undefined {
-  return lineage?.[nodeId]?.columns?.[encodeURI(columnName)]
+  return lineage?.[nodeId]?.columns?.[encodeURI(columnName) as ModelName]
 }

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

@@ -17,6 +17,7 @@ import {
 } from './ModelNode'
 import type { Lineage } from '@/domain/lineage'
 import type { ConnectedNode } from '@/workers/lineage'
+import type { ModelName, ModelFQN } from '@/types/models'
 
 export interface GraphNodeData {
   label: string
@@ -108,7 +109,7 @@ function getEdges(lineage: Record<string, Lineage> = {}): Edge[] {
     })
 
     for (const targetColumnName in targetModel.columns) {
-      const sourceModel = targetModel.columns[targetColumnName]
+      const sourceModel = targetModel.columns[targetColumnName as ModelName]
 
       if (isNil(sourceModel) || isNil(sourceModel.models)) continue
 
@@ -210,7 +211,7 @@ function getNodeMap({
       node.targetPosition = Position.Left
     }
 
-    if (sources.has(node.id)) {
+    if (sources.has(node.id as ModelFQN)) {
       node.sourcePosition = Position.Right
     }
 
@@ -343,7 +344,7 @@ function mergeLineageWithColumns(
       }
 
       // New Column Lineage delivers fresh data, so we can just assign it
-      currentLineageModel.columns[targetColumnNameEncoded] = {
+      currentLineageModel.columns[targetColumnNameEncoded as ModelName] = {
         expression: newLineageModelColumn.expression,
         source: newLineageModelColumn.source,
         models: {},
@@ -353,7 +354,7 @@ function mergeLineageWithColumns(
       if (isObjectEmpty(newLineageModelColumn.models)) continue
 
       const currentLineageModelColumn =
-        currentLineageModel.columns[targetColumnNameEncoded]!
+        currentLineageModel.columns[targetColumnNameEncoded as ModelName]!
       const currentLineageModelColumnModels = currentLineageModelColumn.models
 
       for (const sourceColumnName in newLineageModelColumn.models) {
@@ -371,7 +372,7 @@ function mergeLineageWithColumns(
                   newLineageModelColumnModel,
                 ),
           ),
-        ).map(encodeURI)
+        ).map((uri: string) => encodeURI(uri))
       }
     }
   }
@@ -481,7 +482,7 @@ function getLineageIndex(lineage: Record<string, Lineage> = {}): string {
 
       if (isNotNil(columns)) {
         Object.keys(columns).forEach(columnName => {
-          const column = columns[columnName]
+          const column = columns[columnName as ModelName]
 
           if (isNotNil(column) && isNotNil(column.models)) {
             Object.keys(column.models).forEach(m => allModels.add(m))

vscode/react/src/domain/lineage.ts

@@ -1,6 +1,7 @@
 import { type LineageColumn } from '@/api/client'
+import type { ModelFQN, ModelName } from '@/types/models'
 
 export interface Lineage {
-  models: string[]
-  columns?: Record<string, LineageColumn>
+  models: ModelFQN[]
+  columns?: Record<ModelName, LineageColumn>
 }

vscode/react/src/domain/sqlmesh-model.ts

@@ -8,12 +8,18 @@ import {
   type ModelDefaultCatalog,
   type ModelDefinition,
 } from '@/api/client'
+import type { ModelFQN, ModelName, ModelPath } from '@/types/models'
 import { isArrayNotEmpty } from '@/utils/index'
 import { ModelInitial } from './initial'
 import type { Lineage } from './lineage'
 
-export interface InitialSQLMeshModel extends Model {
-  lineage?: Record<string, Lineage>
+export interface InitialSQLMeshModel
+  extends Omit<Model, 'name' | 'fqn' | 'path' | 'full_path'> {
+  name: ModelName
+  fqn: ModelFQN
+  path: ModelPath
+  full_path: ModelPath
+  lineage?: Record<ModelName, Lineage>
 }
 
 export class ModelSQLMeshModel<
@@ -22,9 +28,10 @@ export class ModelSQLMeshModel<
   _details: ModelDetails = {}
   _detailsIndex: string = ''
 
-  name: string
-  fqn: string
-  path: string
+  name: ModelName
+  fqn: ModelFQN
+  path: ModelPath
+  full_path: ModelPath
   dialect: string
   type: ModelType
   columns: Column[]
@@ -46,10 +53,11 @@ export class ModelSQLMeshModel<
           },
     )
 
-    this.name = encodeURI(this.initial.name)
-    this.fqn = encodeURI(this.initial.fqn)
+    this.name = encodeURI(this.initial.name) as ModelName
+    this.fqn = encodeURI(this.initial.fqn) as ModelFQN
     this.default_catalog = this.initial.default_catalog
-    this.path = this.initial.path
+    this.path = this.initial.path as ModelPath
+    this.full_path = this.initial.full_path as ModelPath
     this.dialect = this.initial.dialect
     this.description = this.initial.description
     this.sql = this.initial.sql
@@ -127,17 +135,21 @@ export class ModelSQLMeshModel<
       } else if (key === 'details') {
         this.details = value as ModelDetails
       } else if (key === 'name') {
-        this.name = encodeURI(value as string)
+        this.name = encodeURI(value as string) as ModelName
       } else if (key === 'fqn') {
-        this.fqn = encodeURI(value as string)
+        this.fqn = encodeURI(value as string) as ModelFQN
       } else if (key === 'type') {
         this.type = value as ModelType
       } else if (key === 'default_catalog') {
         this.default_catalog = value as ModelDefaultCatalog
       } else if (key === 'description') {
         this.description = value as ModelDescription
       } else if (key in this) {
-        this[key as 'path' | 'dialect' | 'sql'] = value as string
+        if (key === 'path' || key === 'full_path') {
+          ;(this as any)[key] = value as ModelPath
+        } else {
+          ;(this as any)[key] = value as string
+        }
       }
     }
   }

vscode/react/src/pages/lineage.tsx

@@ -16,6 +16,7 @@ import type { VSCodeEvent } from '@bus/callbacks'
 import { URI } from 'vscode-uri'
 import type { Model } from '@/api/client'
 import { useRpc } from '@/utils/rpc'
+import type { ModelName, ModelFQN, ModelPath } from '@/types/models'
 
 export function LineagePage() {
   const { emit } = useEventBus()
@@ -198,7 +199,13 @@ export function LineageComponentFromWeb({
   }
 
   const sqlmModel = new ModelSQLMeshModel()
-  sqlmModel.update(model)
+  sqlmModel.update({
+    ...model,
+    name: model.name as ModelName,
+    fqn: model.fqn as ModelFQN,
+    path: model.path as ModelPath,
+    full_path: model.full_path as ModelPath,
+  })
 
   return (
     <div className="h-[100vh] w-[100vw]">

vscode/react/src/types/models.ts

@@ -0,0 +1,6 @@
+export type Brand<T, B extends string> = T & { __brand: B }
+
+export type ModelName = Brand<string, 'ModelName'>
+export type ModelFQN = Brand<string, 'ModelFQN'>
+export type ModelURI = Brand<string, 'ModelURI'>
+export type ModelPath = Brand<string, 'ModelPath'>

vscode/react/src/workers/lineage.ts

@@ -1,5 +1,6 @@
 import { isFalse, isNil, isStringEmptyOrNil, toID } from '@/utils/index'
 import { type Lineage } from '@/domain/lineage'
+import type { ModelFQN } from '@/types/models'
 
 export interface ConnectedNode {
   id?: string
@@ -47,7 +48,7 @@ async function mergeLineageWithModels(
       key = encodeURI(key)
 
       acc[key] = {
-        models: models.map(encodeURI),
+        models: models.map(encodeURI) as ModelFQN[],
         columns: currentLineage?.[key]?.columns ?? undefined,
       }
 
@@ -88,7 +89,7 @@ function getConnectedNodes(
 
   if (isDownstream) {
     models = Object.keys(lineage).filter(key =>
-      lineage[key]!.models.includes(node),
+      lineage[key]!.models.includes(node as ModelFQN),
     )
   } else {
     models = lineage[node]?.models ?? []