Repository: TobikoData/sqlmesh
Pull request #2043 (2024-01-29 20:24:28)

Fix: Slightly improve model autocompletion


Patch:
diff --git a/examples/sushi/macros/macros.py b/examples/sushi/macros/macros.py
index abf0ebb9..f4c2b6a8 100644
--- a/examples/sushi/macros/macros.py
+++ b/examples/sushi/macros/macros.py
@@ -1,8 +1,9 @@
 from macros.utils import between  # type: ignore
+from sqlglot import exp
 
 from sqlmesh import macro
 
 
 @macro()
-def incremental_by_ds(evaluator, column):
+def incremental_by_ds(evaluator, column: exp.Column):
     return between(evaluator, column, evaluator.locals["start_date"], evaluator.locals["end_date"])
diff --git a/sqlmesh/core/engine_adapter/mssql.py b/sqlmesh/core/engine_adapter/mssql.py
index 875b8fa9..1988441e 100644
--- a/sqlmesh/core/engine_adapter/mssql.py
+++ b/sqlmesh/core/engine_adapter/mssql.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import typing as t
 
+import numpy as np
 import pandas as pd
 from pandas.api.types import is_datetime64_any_dtype  # type: ignore
 from sqlglot import exp
@@ -174,7 +175,7 @@ class MSSQLEngineAdapter(
                 columns_to_types_create = columns_to_types.copy()
                 self._convert_df_datetime(df, columns_to_types_create)
                 self.create_table(temp_table, columns_to_types_create)
-                rows: t.List[t.Tuple[t.Any, ...]] = list(df.itertuples(index=False, name=None))  # type: ignore
+                rows: t.List[t.Tuple[t.Any, ...]] = list(df.replace({np.nan: None}).itertuples(index=False, name=None))  # type: ignore
                 conn = self._connection_pool.get()
                 conn.bulk_copy(temp_table.sql(dialect=self.dialect), rows)
             return exp.select(*self._casted_columns(columns_to_types)).from_(temp_table)  # type: ignore
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index 04258443..fc4d5048 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+import inspect
+import logging
 import typing as t
 from enum import Enum
-from functools import reduce
+from functools import reduce, wraps
 from string import Template
 
 import sqlglot
 from jinja2 import Environment
-from sqlglot import Generator, exp
+from sqlglot import Generator, exp, parse_one
 from sqlglot.executor.env import ENV
 from sqlglot.executor.python import Python
 from sqlglot.helper import csv, ensure_collection
@@ -15,6 +17,7 @@ from sqlglot.schema import MappingSchema
 
 from sqlmesh.core.dialect import (
     SQLMESH_MACRO_PREFIX,
+    Dialect,
     MacroDef,
     MacroFunc,
     MacroSQL,
@@ -33,6 +36,8 @@ if t.TYPE_CHECKING:
     from sqlmesh.core.engine_adapter import EngineAdapter
     from sqlmesh.core.snapshot import Snapshot
 
+logger = logging.getLogger(__name__)
+
 
 class RuntimeStage(Enum):
     LOADING = "loading"
@@ -349,6 +354,75 @@ class MacroEvaluator:
             )
         return self.locals["engine_adapter"]
 
+    def _coerce(self, expr: exp.Expression, typ: t.Any, strict: bool = False) -> t.Any:
+        """Coerces the given expression to the specified type on a best-effort basis."""
+        base_err_msg = f"Failed to coerce expression '{expr}' to type '{typ}'."
+        try:
+            if typ is None:
+                return expr
+            base = t.get_origin(typ) or typ
+            # We need to handle t.Union first since we cannot use isinstance with it
+            if base is t.Union:
+                for branch in t.get_args(typ):
+                    try:
+                        return self._coerce(expr, branch, True)
+                    except Exception:
+                        pass
+                raise SQLMeshError(base_err_msg)
+            if isinstance(expr, base):
+                return expr
+            if issubclass(base, exp.Expression):
+                d = Dialect.get_or_raise(self.dialect)
+                into = base if base in d.parser().EXPRESSION_PARSERS else None
+                if into is None:
+                    if isinstance(expr, exp.Literal):
+                        coerced = parse_one(expr.this)
+                    else:
+                        raise SQLMeshError(
+                            f"{base_err_msg} Coercion to {base} requires a literal expression."
+                        )
+                else:
+                    coerced = parse_one(
+                        expr.this if isinstance(expr, exp.Literal) else expr.sql(), into=into
+                    )
+                if isinstance(coerced, base):
+                    return coerced
+                else:
+                    raise SQLMeshError(base_err_msg)
+            if base in (int, float, str) and isinstance(expr, exp.Literal):
+                return base(expr.this)
+            if base is bool and isinstance(expr, exp.Boolean):
+                return expr.this
+            if base is str and isinstance(expr, exp.Expression):
+                return expr.sql(self.dialect)
+            if base is tuple and isinstance(expr, (exp.Tuple, exp.Array)):
+                generic = t.get_args(typ)
+                if not generic:
+                    return tuple(expr.expressions)
+                if generic[-1] is ...:
+                    return tuple(self._coerce(expr, generic[0]) for expr in expr.expressions)
+                elif len(generic) == len(expr.expressions):
+                    return tuple(
+                        self._coerce(expr, generic[i]) for i, expr in enumerate(expr.expressions)
+                    )
+                raise SQLMeshError(f"{base_err_msg} Expected {len(generic)} items.")
+            if base is list and isinstance(expr, (exp.Array, exp.Tuple)):
+                generic = t.get_args(typ)
+                if not generic:
+                    return expr.expressions
+                return [self._coerce(expr, generic[0]) for expr in expr.expressions]
+            raise SQLMeshError(base_err_msg)
+        except Exception:
+            if strict:
+                raise
+            logger.warning(
+                "Coercion of expression '%s' to type '%s' failed. Using non coerced expression.",
+                expr,
+                typ,
+                exc_info=True,
+            )
+            return expr
+
 
 class macro(registry_decorator):
     """Specifies a function is a macro and registers it the global MACROS registry.
@@ -372,7 +446,31 @@ class macro(registry_decorator):
     def __call__(
         self, func: t.Callable[..., DECORATOR_RETURN_TYPE]
     ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
-        wrapper = super().__call__(func)
+        @wraps(func)
+        def _typed_func(
+            evaluator: MacroEvaluator, *args_: t.Any, **kwargs_: t.Any
+        ) -> DECORATOR_RETURN_TYPE:
+            spec = inspect.getfullargspec(func)
+            annotations = t.get_type_hints(func)
+            kwargs = inspect.getcallargs(func, evaluator, *args_, **kwargs_)
+            for param, value in kwargs.items():
+                coercible_type = annotations.get(param)
+                if not coercible_type:
+                    continue
+                kwargs[param] = evaluator._coerce(value, coercible_type)
+            args = [kwargs.pop(k) for k in spec.args if k in kwargs]
+            if spec.varargs:
+                args.extend(kwargs.pop(spec.varargs, []))
+            return func(*args, **kwargs)
+
+        try:
+            annotated = any(t.get_type_hints(func).keys() - {"return"})
+        except TypeError:
+            annotated = False
+
+        wrapper = super().__call__(
+            func if not annotated else t.cast(t.Callable[..., DECORATOR_RETURN_TYPE], _typed_func)
+        )
 
         # This is useful to identify macros at runtime
         setattr(wrapper, "__sqlmesh_macro__", True)
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index ee0d9ae8..32a2efa5 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -1170,7 +1170,9 @@ class SeedModel(_SqlBasedModel):
             for column in date_or_time_columns:
                 df[column] = pd.to_datetime(df[column])
             df[bool_columns] = df[bool_columns].apply(lambda i: str_to_bool(str(i)))
-            df[string_columns] = df[string_columns].astype(str)
+            df[string_columns] = df[string_columns].mask(
+                cond=lambda x: x.notna(), other=df[string_columns].astype(str)  # type: ignore
+            )
             yield df
 
     def text_diff(self, other: Node) -> str:
diff --git a/web/client/src/library/components/editor/Editor.tsx b/web/client/src/library/components/editor/Editor.tsx
index a99bf86f..31981253 100644
--- a/web/client/src/library/components/editor/Editor.tsx
+++ b/web/client/src/library/components/editor/Editor.tsx
@@ -20,6 +20,8 @@ import { type KeyBinding } from '@codemirror/view'
 import { useStoreContext } from '@context/context'
 import { useIDE } from '~/library/pages/ide/context'
 import { ModelDirectory } from '@models/directory'
+import { type ModelSQLMeshModel } from '@models/sqlmesh-model'
+import { type Column } from '@api/client'
 
 function Editor(): JSX.Element {
   const tab = useStoreEditor(s => s.tab)
@@ -70,15 +72,23 @@ function EditorMain({ tab }: { tab: EditorTab }): JSX.Element {
 
   const { setManuallySelectedColumn } = useLineageFlow()
 
+  const handleModelClick = useCallback(
+    function (model: ModelSQLMeshModel) {
+      setSelectedFile(files.get(model.path))
+    },
+    [files],
+  )
+  const handleModelColumn = useCallback(function (
+    model: ModelSQLMeshModel,
+    column: Column,
+  ) {
+    setManuallySelectedColumn([model, column])
+  }, [])
   const defaultKeymapsEditorTab = useDefaultKeymapsEditorTab()
   const modelExtensions = useSQLMeshModelExtensions(
     tab.file.path,
-    model => {
-      setSelectedFile(files.get(model.path))
-    },
-    (model, column) => {
-      setManuallySelectedColumn([model, column])
-    },
+    handleModelClick,
+    handleModelColumn,
   )
 
   const [isOpenInspector, setIsOpenInspector] = useState(false)
@@ -258,6 +268,7 @@ function EditorMain({ tab }: { tab: EditorTab }): JSX.Element {
                 dialect={tab.dialect}
                 keymaps={customSQLKeymaps}
                 content={tab.file.content}
+                extensions={modelExtensions}
                 onChange={updateFileContent}
               />
             )}
diff --git a/web/client/src/library/components/editor/EditorCode.tsx b/web/client/src/library/components/editor/EditorCode.tsx
index 81cf9334..c9609d25 100644
--- a/web/client/src/library/components/editor/EditorCode.tsx
+++ b/web/client/src/library/components/editor/EditorCode.tsx
@@ -109,13 +109,40 @@ function CodeEditorDefault({
       ),
     [keymaps],
   )
+  const allModels = useMemo(() => Array.from(models.values()), [models])
+  const allModelsNames = useMemo(
+    () =>
+      allModels.map(model => ({
+        label: model.name,
+        type: 'model',
+      })),
+    [allModels],
+  )
+  const allColumnsNames = useMemo(
+    () =>
+      Array.from(
+        new Set(
+          allModels.flatMap(model => model.columns.map(column => column.name)),
+        ),
+      ).map(name => ({
+        label: name,
+        type: 'column',
+      })),
+    [allModels],
+  )
 
   const extensionsAll = useMemo(() => {
     return [
       ...extensionsDefault,
       extensions,
       type === EnumFileExtensions.SQL &&
-        SQLMeshDialect(models, dialectOptions, dialectsTitles),
+        SQLMeshDialect(
+          models,
+          allModelsNames,
+          allColumnsNames,
+          dialectOptions,
+          dialectsTitles,
+        ),
       extensionKeymap,
     ]
       .filter(Boolean)
@@ -128,6 +155,8 @@ function CodeEditorDefault({
     extensionsDefault,
     extensions,
     extensionKeymap,
+    allModelsNames,
+    allColumnsNames,
   ])
 
   useEffect(() => {
diff --git a/web/client/src/library/components/editor/EditorPreview.tsx b/web/client/src/library/components/editor/EditorPreview.tsx
index 9118e430..94edfbd1 100644
--- a/web/client/src/library/components/editor/EditorPreview.tsx
+++ b/web/client/src/library/components/editor/EditorPreview.tsx
@@ -61,6 +61,9 @@ export default function EditorPreview({
   const model = models.get(tab.file.path)
   const showLineage =
     isFalse(tab.file.isEmpty) && isNotNil(model) && isModel(tab.file.path)
+  const showErrors =
+    errors.size > 0 &&
+    ([previewTable, previewDiff].some(isNotNil) || showLineage)
 
   const tabs: string[] = useMemo(
     () =>
@@ -69,13 +72,21 @@ export default function EditorPreview({
         isNotNil(previewQuery) && EnumEditorPreviewTabs.Query,
         showLineage && EnumEditorPreviewTabs.Lineage,
         isNotNil(previewDiff) && EnumEditorPreviewTabs.Diff,
-        errors.size > 0 && EnumEditorPreviewTabs.Errors,
+        showErrors && EnumEditorPreviewTabs.Errors,
       ].filter(Boolean) as string[],
-    [tab.id, previewTable, previewQuery, previewDiff, showLineage, errors],
+    [
+      tab.id,
+      previewTable,
+      previewQuery,
+      previewDiff,
+      showLineage,
+      errors,
+      showErrors,
+    ],
   )
 
   useEffect(() => {
-    if (errors.size > 0) {
+    if (showErrors) {
       setActiveTabIndex(tabs.indexOf(EnumEditorPreviewTabs.Errors))
     } else if (isNotNil(previewDiff)) {
       setActiveTabIndex(tabs.indexOf(EnumEditorPreviewTabs.Diff))
@@ -84,7 +95,15 @@ export default function EditorPreview({
     } else if (showLineage) {
       setActiveTabIndex(tabs.indexOf(EnumEditorPreviewTabs.Lineage))
     }
-  }, [tabs, previewTable, previewQuery, previewDiff, showLineage, errors])
+  }, [
+    tabs,
+    previewTable,
+    previewQuery,
+    previewDiff,
+    showLineage,
+    errors,
+    showErrors,
+  ])
 
   return (
     <div
@@ -182,7 +201,7 @@ export default function EditorPreview({
                 />
               </Tab.Panel>
             )}
-            {errors.size > 0 && (
+            {showErrors && (
               <Tab.Panel
                 unmount={false}
                 className="w-full h-full ring-white ring-opacity-60 ring-offset-2 ring-offset-blue-400 focus:outline-none focus:ring-2 py-2"
diff --git a/web/client/src/library/components/editor/extensions/SQLMeshDialect.ts b/web/client/src/library/components/editor/extensions/SQLMeshDialect.ts
index 475e3ea9..7c37f964 100644
--- a/web/client/src/library/components/editor/extensions/SQLMeshDialect.ts
+++ b/web/client/src/library/components/editor/extensions/SQLMeshDialect.ts
@@ -1,3 +1,4 @@
+import { type Model } from '@api/client'
 import {
   completeFromList,
   type CompletionContext,
@@ -5,7 +6,6 @@ import {
 } from '@codemirror/autocomplete'
 import { keywordCompletionSource, SQLDialect } from '@codemirror/lang-sql'
 import { LanguageSupport } from '@codemirror/language'
-import { type Model } from '~/api/client'
 import { isArrayEmpty, isNil, isNotNil, isStringEmptyOrNil } from '~/utils'
 import { sqlglotWorker } from '~/workers'
 
@@ -14,12 +14,16 @@ const WHITE_SPACE = ' '
 
 export type ExtensionSQLMeshDialect = (
   models: Map<string, Model>,
+  allModelsNames: Array<{ label: string; type: string }>,
+  allColumnsNames: Array<{ label: string; type: string }>,
   options?: { types: string; keywords: string },
   dialects?: string[],
 ) => LanguageSupport
 
 export const SQLMeshDialect: ExtensionSQLMeshDialect = function SQLMeshDialect(
   models,
+  allModelsNames,
+  allColumnsNames,
   options = { types: '', keywords: '' },
   dialects,
 ): LanguageSupport {
@@ -34,20 +38,6 @@ export const SQLMeshDialect: ExtensionSQLMeshDialect = function SQLMeshDialect(
     keywords: (SQLKeywords + WHITE_SPACE + SQLMeshKeywords).toLowerCase(),
     types: (SQLTypes + WHITE_SPACE + SQLMeshTypes).toLowerCase(),
   })
-  const allModels = Array.from(new Set(models.values()))
-  const modelColumns = Array.from(
-    new Set(
-      allModels.map(model => model.columns.map(column => column.name)).flat(),
-    ),
-  )
-  const modelNames: Completion[] = allModels.map(model => ({
-    label: model.name,
-    type: 'model',
-  }))
-  const columnNames: Completion[] = modelColumns.map(name => ({
-    label: name,
-    type: 'column',
-  }))
 
   SQLMeshDialectCleanUp()
 
@@ -67,11 +57,11 @@ export const SQLMeshDialect: ExtensionSQLMeshDialect = function SQLMeshDialect(
         const dot = ctx.matchBefore(/[A-Za-z0-9_.]*\.(\w+)?\s*$/i)
 
         if (isNotNil(dot)) {
-          let options = columnNames
+          let options = allColumnsNames
           const blocks = dot.text.split('.')
           const text = blocks.pop()
           const maybeModelName = blocks.join('.')
-          const maybeModel = models.get(maybeModelName)
+          const maybeModel = models.get(maybeModelName) ?? models.get(dot.text)
 
           if (isNotNil(maybeModel)) {
             options = maybeModel.columns.map(column => ({
@@ -81,7 +71,7 @@ export const SQLMeshDialect: ExtensionSQLMeshDialect = function SQLMeshDialect(
           }
 
           return {
-            from: isStringEmptyOrNil(text) ? dot.to : dot.to - text.length,
+            from: isStringEmptyOrNil(text) ? dot.to : dot.to - dot.text.length,
             to: dot.to,
             options,
           }
@@ -94,9 +84,6 @@ export const SQLMeshDialect: ExtensionSQLMeshDialect = function SQLMeshDialect(
         const keywordKind = matchWordWithSpacesAfter(ctx, 'kind')
         const keywordDialect = matchWordWithSpacesAfter(ctx, 'dialect')
         const keywordModel = matchWordWithSpacesAfter(ctx, 'model')
-        const keywordFrom = matchWordWithSpacesAfter(ctx, 'from')
-        const keywordJoin = matchWordWithSpacesAfter(ctx, 'join')
-        const keywordSelect = matchWordWithSpacesAfter(ctx, 'select')
 
         const text = ctx.state.doc.toJSON().join('\n')
         const matchModels = text.match(/MODEL \(([\s\S]+?)\);/g) ?? []
@@ -121,27 +108,48 @@ export const SQLMeshDialect: ExtensionSQLMeshDialect = function SQLMeshDialect(
           }
 
           return completeFromList(suggestions ?? [])(ctx)
-        } else {
-          let suggestions: Completion[] = []
-
-          if (isNotNil(keywordModel) && isArrayEmpty(matchModels)) {
-            suggestions = [
-              {
-                label: 'model',
-                type: 'keyword',
-                apply: 'MODEL (\n\r);',
-              },
-            ]
-          } else if (isNotNil(keywordSelect)) {
-            suggestions = columnNames
-          } else if (isNotNil(keywordFrom) || isNotNil(keywordJoin)) {
-            suggestions = modelNames
-          }
+        }
+
+        let suggestions: Completion[] = []
+
+        if (isNotNil(keywordModel) && isArrayEmpty(matchModels)) {
+          suggestions = [
+            {
+              label: 'model',
+              type: 'keyword',
+              apply: 'MODEL (\n\r);',
+            },
+          ]
+        }
+
+        if (maybeSuggestion(allModelsNames, word.text)) {
+          suggestions.push(...allModelsNames)
+        }
 
-          return isArrayEmpty(suggestions)
-            ? keywordCompletionSource(lang)(ctx)
-            : completeFromList(suggestions)(ctx)
+        if (maybeSuggestion(allColumnsNames, word.text)) {
+          suggestions.push(...allColumnsNames)
         }
+
+        const wordLastChar = word.text.slice(-1)
+        const isUpperCase = wordLastChar === wordLastChar.toUpperCase()
+
+        if (isUpperCase) {
+          suggestions = suggestions.map(suggestion => ({
+            ...suggestion,
+            label: suggestion.label.toUpperCase(),
+          }))
+        }
+
+        const source = isArrayEmpty(suggestions)
+          ? keywordCompletionSource(lang, isUpperCase)(ctx)
+          : {
+              from: word.from,
+              to: word.to,
+              options: suggestions,
+              validFor: /^\w*$/,
+            }
+
+        return source
       },
     }),
   ])
@@ -241,3 +249,10 @@ function matchWordWithSpacesAfter(
 }> {
   return ctx.matchBefore(new RegExp(`${word}\\s*(?=\\S)([^ ]+)`, 'i'))
 }
+
+function maybeSuggestion(suggestions: Completion[], word: string): boolean {
+  return (
+    word.length > 1 &&
+    suggestions.some(suggestion => suggestion.label.includes(word))
+  )
+}
diff --git a/web/client/src/library/components/editor/extensions/SQLMeshModel.ts b/web/client/src/library/components/editor/extensions/SQLMeshModel.ts
index 7d45edbd..389801cc 100644
--- a/web/client/src/library/components/editor/extensions/SQLMeshModel.ts
+++ b/web/client/src/library/components/editor/extensions/SQLMeshModel.ts
@@ -11,9 +11,9 @@ import { getDecorations } from './help'
 
 export default function SQLMeshModel(
   models: Map<string, ModelSQLMeshModel>,
-  model: ModelSQLMeshModel,
   columns: Set<string>,
   isActionMode: boolean,
+  model?: ModelSQLMeshModel,
 ): Extension {
   return ViewPlugin.fromClass(
     class SQLMeshModelView {
@@ -22,9 +22,9 @@ export default function SQLMeshModel(
         this.decorations = getDecorations(
           models,
           view,
-          model,
           columns,
           isActionMode,
+          model,
         )
       }
 
@@ -32,9 +32,9 @@ export default function SQLMeshModel(
         this.decorations = getDecorations(
           models,
           viewUpdate.view,
-          model,
           columns,
           isActionMode,
+          model,
         )
       }
     },
diff --git a/web/client/src/library/components/editor/extensions/help.ts b/web/client/src/library/components/editor/extensions/help.ts
index 0a08ecb3..66ba3f94 100644
--- a/web/client/src/library/components/editor/extensions/help.ts
+++ b/web/client/src/library/components/editor/extensions/help.ts
@@ -52,14 +52,15 @@ export function findColumn(
 export function getDecorations(
   models: Map<string, ModelSQLMeshModel>,
   view: EditorView,
-  model: ModelSQLMeshModel,
   columns: Set<string>,
   isActionMode: boolean,
+  model?: ModelSQLMeshModel,
 ): DecorationSet {
   const ranges: any = []
   const columnNames = Array.from(columns)
-  const modelNames = Array.from(new Set(models.values())).map(m => m.name)
-  const modelColumns = model.columns.map(c => c.name.toLowerCase())
+  const allModels = Array.from(new Set(models.values()))
+  const modelNames = allModels.map(m => m.name)
+  const modelColumns = model?.columns.map(c => c.name.toLowerCase()) ?? []
   const validLeftCharColumn = new Set(['.', '(', '[', ' ', '\n'])
   const validRightCharColumn = new Set([':', ')', ',', ']', ' ', '\n'])
 
@@ -73,14 +74,14 @@ export function getDecorations(
           createMarkDecorationModel({
             model: name,
             isActionMode,
-            isActiveModel: model.name === name,
+            isActiveModel: model?.name === name,
           }).range(from, to),
         )
       },
     )
 
     getMarkDecorations(
-      columnNames,
+      isNil(model) ? columnNames : modelColumns,
       view.state.doc.sliceString(range.from, range.to),
       range,
       ({ from, to, key }) => {
diff --git a/web/client/src/library/components/editor/hooks.ts b/web/client/src/library/components/editor/hooks.ts
index 4db8cb4a..7e6508ad 100644
--- a/web/client/src/library/components/editor/hooks.ts
+++ b/web/client/src/library/components/editor/hooks.ts
@@ -4,7 +4,7 @@ import { useLineageFlow } from '@components/graph/context'
 import { useStoreEditor } from '@context/editor'
 import { type ModelSQLMeshModel } from '@models/sqlmesh-model'
 import { type Column } from '~/api/client'
-import { isFalse, isNil, isNotNil } from '@utils/index'
+import { isFalse, isNil, isNotNil, isObjectEmpty } from '@utils/index'
 import { useMemo, useState } from 'react'
 import { events, HoverTooltip, SQLMeshModel } from './extensions'
 import { findModel, findColumn } from './extensions/help'
@@ -69,14 +69,20 @@ function useSQLMeshModelExtensions(
   const [isActionMode, setIsActionMode] = useState(false)
 
   const extensions = useMemo(() => {
-    const columns = isNil(lineage)
-      ? new Set<string>()
-      : new Set(
-          Object.keys(lineage)
-            .map(modelName => models.get(modelName)?.columns.map(c => c.name))
-            .flat()
-            .filter(Boolean) as string[],
-        )
+    const columns =
+      isNil(lineage) || isObjectEmpty(lineage)
+        ? new Set<string>(
+            Array.from(new Set(models.values())).flatMap(m =>
+              m.columns.map(c => c.name),
+            ),
+          )
+        : new Set(
+            Object.keys(lineage)
+              .flatMap(
+                modelName => models.get(modelName)?.columns.map(c => c.name),
+              )
+              .filter(Boolean) as string[],
+          )
 
     function handleEventModelClick(event: MouseEvent): void {
       if (event.metaKey) {
@@ -116,7 +122,7 @@ function useSQLMeshModelExtensions(
       }),
       isNotNil(handleModelClick) && events({ click: handleEventModelClick }),
       isNotNil(handleModelColumn) && events({ click: handleEventlColumnClick }),
-      isNotNil(model) && SQLMeshModel(models, model, columns, isActionMode),
+      SQLMeshModel(models, columns, isActionMode, model),
     ].filter(Boolean) as Extension[]
   }, [handleModelClick, handleModelColumn, model, models, files, isActionMode])
 


Test patch:
diff --git a/tests/core/engine_adapter/test_integration.py b/tests/core/engine_adapter/test_integration.py
index 6a149bee..ee377034 100644
--- a/tests/core/engine_adapter/test_integration.py
+++ b/tests/core/engine_adapter/test_integration.py
@@ -7,6 +7,7 @@ import sys
 import typing as t
 from datetime import timedelta
 
+import numpy as np
 import pandas as pd
 import pytest
 from sqlglot import exp, parse_one
@@ -174,7 +175,7 @@ class TestContext:
     def get_current_data(self, table: exp.Table) -> pd.DataFrame:
         df = self.engine_adapter.fetchdf(exp.select("*").from_(table), quote_identifiers=True)
         if self.dialect == "snowflake" and "id" in df.columns:
-            df["id"] = df["id"].astype("int")
+            df["id"] = df["id"].apply(lambda x: x if pd.isna(x) else int(x))
         return df
 
     def compare_with_current(self, table: exp.Table, expected: pd.DataFrame) -> None:
@@ -820,6 +821,34 @@ def test_drop_schema(ctx: TestContext):
     assert len(results.views) == 0
 
 
+def test_nan_roundtrip(ctx: TestContext):
+    if ctx.test_type != "df":
+        pytest.skip("NaN roundtrip test only relevant for dataframes.")
+    ctx.engine_adapter.DEFAULT_BATCH_SIZE = sys.maxsize
+    ctx.init()
+    table = ctx.table("test_table")
+    # Initial Load
+    input_data = pd.DataFrame(
+        [
+            {"id": 1, "ds": "2022-01-01"},
+            {"id": 2, "ds": "2022-01-02"},
+            {"id": np.nan, "ds": np.nan},
+        ]
+    )
+    ctx.engine_adapter.create_table(table, ctx.columns_to_types)
+    ctx.engine_adapter.replace_query(
+        table,
+        ctx.input_data(input_data),
+        columns_to_types=ctx.columns_to_types,
+    )
+    results = ctx.get_metadata_results()
+    assert not results.views
+    assert not results.materialized_views
+    assert len(results.tables) == len(results.non_temp_tables) == 1
+    assert results.non_temp_tables[0] == table.name
+    ctx.compare_with_current(table, input_data)
+
+
 def test_replace_query(ctx: TestContext):
     ctx.engine_adapter.DEFAULT_BATCH_SIZE = sys.maxsize
     ctx.init()
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index 74508c6b..cadf8db3 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -27,6 +27,41 @@ def macro_evaluator() -> MacroEvaluator:
     def noop(evaluator: MacroEvaluator):
         return None
 
+    @macro()
+    def bitshift_square(evaluator: MacroEvaluator, x: int, y: int) -> int:
+        return (x >> y) ** 2
+
+    @macro()
+    def prefix_db(evaluator: MacroEvaluator, table: exp.Table, prefix: str) -> exp.Table:
+        table.set("db", prefix + table.db)
+        return table
+
+    @macro()
+    def repeated(evaluator: MacroEvaluator, expr: str, times: int = 2, multi: bool = False):
+        if multi is True:
+            return (expr,) * times
+        return expr * times
+
+    @macro()
+    def split(evaluator: MacroEvaluator, string: str, sep: str = ","):
+        return string.split(sep)
+
+    @macro()
+    def cte_tag_name(evaluator: MacroEvaluator, with_: exp.Select):
+        for cte in with_.find_all(exp.CTE):
+            name = cte.alias_or_name
+            for query in cte.find_all(exp.Select):
+                query.select(exp.Literal.string(name).as_("source"), copy=False)
+        return with_
+
+    @macro()
+    def suffix_idents(evaluator: MacroEvaluator, items: t.List[str], suffix: str):
+        return [item + suffix for item in items]
+
+    @macro()
+    def suffix_idents_2(evaluator: MacroEvaluator, items: t.Tuple[str, ...], suffix: str):
+        return [item + suffix for item in items]
+
     return MacroEvaluator(
         "hive",
         {"test": Executable(name="test", payload=f"def test(_):\n    return 'test'")},
@@ -292,12 +327,131 @@ def test_ast_correctness(macro_evaluator):
             "SELECT d",
             {},
         ),
+        (
+            """@BITSHIFT_SQUARE(50, '3')""",
+            "36",
+            {},
+        ),
+        (
+            """@PREFIX_DB(my.schema.table, 'dev_')""",
+            "my.dev_schema.table",
+            {},
+        ),
+        (
+            """select @REPEATED(test, 3)""",
+            "SELECT testtesttest",
+            {},
+        ),
+        (
+            """select @REPEATED(test, 3, true)""",
+            "SELECT test, test, test",
+            {},
+        ),
+        (
+            """select @SPLIT('a,b,c')""",
+            "SELECT a, b, c",
+            {},
+        ),
+        (
+            """@CTE_TAG_NAME(WITH step1 AS (SELECT 1) SELECT * FROM step1)""",
+            "WITH step1 AS (SELECT 1, 'step1' AS source) SELECT * FROM step1",
+            {},
+        ),
+        (
+            """@CTE_TAG_NAME('WITH step1 AS (SELECT 1) SELECT * FROM step1')""",
+            "WITH step1 AS (SELECT 1, 'step1' AS source) SELECT * FROM step1",
+            {},
+        ),
+        (
+            """SELECT @SUFFIX_IDENTS(['a', 'b', 'c'], 'z')""",
+            "SELECT az, bz, cz",
+            {},
+        ),
+        (
+            """SELECT @SUFFIX_IDENTS_2(['a', 'b', 'c'], 'z')""",
+            "SELECT az, bz, cz",
+            {},
+        ),
     ],
 )
-def test_macro_functions(macro_evaluator, assert_exp_eq, sql, expected, args):
+def test_macro_functions(macro_evaluator: MacroEvaluator, assert_exp_eq, sql, expected, args):
     macro_evaluator.locals = args or {}
     assert_exp_eq(macro_evaluator.transform(parse_one(sql)), expected)
 
 
-def test_macro_returns_none(macro_evaluator):
+def test_macro_returns_none(macro_evaluator: MacroEvaluator):
     assert macro_evaluator.transform(parse_one("@NOOP()")) is None
+
+
+def test_macro_coercion(macro_evaluator: MacroEvaluator, assert_exp_eq):
+    coerce = macro_evaluator._coerce
+    assert coerce(exp.Literal.number(1), int) == 1
+    assert coerce(exp.Literal.number(1.1), float) == 1.1
+    assert coerce(exp.Literal.string("Hi mom"), str) == "Hi mom"
+    assert coerce(exp.true(), bool) is True
+
+    # Coercing a string literal to a column should return a column with the same name
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Column), exp.column("order"))
+    # Not possible to coerce this string literal Cast to an exp.Column node -- so it should just return the input
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), exp.Column), exp.Literal.string("order::date")
+    )
+    # This however, is correctly coercible since it's a cast
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), exp.Cast), exp.cast(exp.column("order"), "DATE")
+    )
+
+    # Here we resolve ambiguity via the user type hint
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Identifier), exp.to_identifier("order"))
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Table), exp.table_("order"))
+
+    # Resolve a union type hint by choosing the first one that works
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), t.Union[exp.Column, exp.Cast]),
+        exp.cast(exp.column("order"), "DATE"),
+    )
+
+    # Simply ask for a string, and always get a string
+    assert coerce(exp.column("order"), str) == "order"
+    assert (
+        coerce(parse_one("SELECT x FROM UNNEST(y) AS x WHERE 1 = 1"), str)
+        == "SELECT x FROM UNNEST(y) AS x WHERE 1 = 1"
+    )
+
+    # From a string literal to a Select should parse the string literal, and the inverse operation works as well
+    assert_exp_eq(
+        coerce(exp.Literal.string("SELECT 1 FROM a"), exp.Select), parse_one("SELECT 1 FROM a")
+    )
+    assert coerce(parse_one("SELECT 1 FROM a"), str) == "SELECT 1 FROM a"
+
+    # Get a list of exp directly instead of an exp.Array
+    assert coerce(parse_one("[1, 2, 3]"), list) == [
+        exp.Literal.number(1),
+        exp.Literal.number(2),
+        exp.Literal.number(3),
+    ]
+
+    # Generics work as well, recursively resolving inner types
+    assert coerce(parse_one("[1, 2, 3]"), t.List[int]) == [1, 2, 3]
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, int, float]) == (1, 2, 3.0)
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, ...]) == (1, 2, 3)
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, str, float]) == (1, "2", 3.0)
+    assert coerce(
+        parse_one("[1, 2, [3]]"), t.Tuple[int, str, t.Union[float, t.Tuple[float, ...]]]
+    ) == (1, "2", (3.0,))
+
+    # Using exp.Expression will always return the input expression
+    assert coerce(parse_one("order", into=exp.Column), exp.Expression) == exp.column("order")
+    assert coerce(exp.Literal.string("OK"), exp.Expression) == exp.Literal.string("OK")
+
+    # Strict flag allows raising errors and is used when recursively coercing expressions
+    # otherwise, in general, we want to be lenient and just warn the user when something is not possible
+    with pytest.raises(SQLMeshError):
+        coerce(exp.Literal.string("order"), exp.Select, strict=True)
+
+    with pytest.raises(SQLMeshError):
+        _ = coerce(
+            exp.Literal.string("order::date"),
+            t.Union[exp.Column, exp.Identifier],
+            strict=True,
+        )
diff --git a/tests/core/test_selector.py b/tests/core/test_selector.py
index b3740b0e..6d4fc537 100644
--- a/tests/core/test_selector.py
+++ b/tests/core/test_selector.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+from pathlib import Path
 from unittest.mock import call
 
 import pytest
@@ -373,6 +374,51 @@ def test_expand_model_selections(
     assert selector.expand_model_selections(selections) == output
 
 
+@pytest.mark.parametrize(
+    "expressions, expected_fqns",
+    [
+        (["git:main"], {'"test_model_a"', '"test_model_c"'}),
+        (["git:main & +*model_c"], {'"test_model_c"'}),
+    ],
+)
+def test_expand_git_selection(
+    mocker: MockerFixture, expressions: t.List[str], expected_fqns: t.Set[str]
+):
+    models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
+
+    model_a = SqlModel(name="test_model_a", query=d.parse_one("SELECT 1 AS a"))
+    model_a._path = Path("/path/to/test_model_a.sql")
+    models[model_a.fqn] = model_a
+
+    model_b = SqlModel(name="test_model_b", query=d.parse_one("SELECT 2 AS b"))
+    model_b._path = Path("/path/to/test_model_b.sql")
+    models[model_b.fqn] = model_b
+
+    model_c = SqlModel(name="test_model_c", query=d.parse_one("SELECT 3 AS c"))
+    model_c._path = Path("/path/to/test_model_c.sql")
+    models[model_c.fqn] = model_c
+
+    model_d = SqlModel(
+        name="test_model_d",
+        query=d.parse_one("SELECT c FROM test_model_c"),
+        depends_on={"test_model_c"},
+    )
+    model_d._path = Path("/path/to/test_model_d.sql")
+    models[model_d.fqn] = model_d
+
+    git_client_mock = mocker.Mock()
+    git_client_mock.list_untracked_files.return_value = []
+    git_client_mock.list_changed_files.return_value = [model_a._path, model_c._path]
+
+    selector = Selector(mocker.Mock(), models)
+    selector._git_client = git_client_mock
+
+    assert selector.expand_model_selections(expressions) == expected_fqns
+
+    git_client_mock.list_changed_files.assert_called_once_with(target_branch="main")
+    git_client_mock.list_untracked_files.assert_called_once()
+
+
 def test_select_models_with_external_parent(mocker: MockerFixture):
     default_catalog = "test_catalog"
     added_model = SqlModel(