feat: add inlay type hints for sqlmesh sql models (#4641)

Christopher Giroir · web-flow · commit ff87b036b2f3 · 2025-06-05T11:28:33.000-07:00
diff --git a/sqlmesh/lsp/hints.py b/sqlmesh/lsp/hints.py
@@ -0,0 +1,136 @@
+"""Type hinting on SQLMesh models"""
+
+import typing as t
+
+from lsprotocol import types
+
+from sqlglot import exp
+from sqlglot.expressions import Expression
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.uri import URI
+
+
+def get_hints(
+    lsp_context: LSPContext,
+    document_uri: URI,
+    start_line: int,
+    end_line: int,
+) -> t.List[types.InlayHint]:
+    """
+    Get type hints for certain lines in a document
+
+    Args:
+        lint_context: The LSP context
+        document_uri: The URI of the document
+        start_line: the starting line to get hints for
+        end_line: the ending line to get hints for
+
+    Returns:
+        A list of hints to apply to the document
+    """
+    path = document_uri.to_path()
+    if path.suffix != ".sql":
+        return []
+
+    if path not in lsp_context.map:
+        return []
+
+    file_info = lsp_context.map[path]
+
+    # Process based on whether it's a model or standalone audit
+    if not isinstance(file_info, ModelTarget):
+        return []
+
+    # It's a model
+    model = lsp_context.context.get_model(
+        model_or_snapshot=file_info.names[0], raise_if_missing=False
+    )
+    if not isinstance(model, SqlModel):
+        return []
+
+    query = model.query
+    dialect = model.dialect
+    columns_to_types = model.columns_to_types or {}
+
+    return _get_type_hints_for_model_from_query(
+        query, dialect, columns_to_types, start_line, end_line
+    )
+
+
+def _get_type_hints_for_select(
+    expression: exp.Expression,
+    dialect: str,
+    columns_to_types: t.Dict[str, exp.DataType],
+    start_line: int,
+    end_line: int,
+) -> t.List[types.InlayHint]:
+    hints: t.List[types.InlayHint] = []
+
+    for select_exp in expression.expressions:
+        if isinstance(select_exp, exp.Alias):
+            if isinstance(select_exp.this, exp.Cast):
+                continue
+
+            meta = select_exp.args["alias"].meta
+
+        elif isinstance(select_exp, exp.Column):
+            meta = select_exp.parts[-1].meta
+        else:
+            continue
+
+        if "line" not in meta or "col" not in meta:
+            continue
+
+        line = meta["line"]
+        col = meta["col"]
+
+        # Lines from sqlglot are 1 based
+        line -= 1
+
+        if line < start_line or line > end_line:
+            continue
+
+        name = select_exp.alias_or_name
+        data_type = columns_to_types.get(name)
+
+        if not data_type or data_type.is_type(exp.DataType.Type.UNKNOWN):
+            continue
+
+        type_label = data_type.sql(dialect)
+        hints.append(
+            types.InlayHint(
+                label=f"::{type_label}",
+                kind=types.InlayHintKind.Type,
+                padding_left=False,
+                padding_right=True,
+                position=types.Position(line=line, character=col),
+            )
+        )
+
+    return hints
+
+
+def _get_type_hints_for_model_from_query(
+    query: Expression,
+    dialect: str,
+    columns_to_types: t.Dict[str, exp.DataType],
+    start_line: int,
+    end_line: int,
+) -> t.List[types.InlayHint]:
+    hints: t.List[types.InlayHint] = []
+    try:
+        query = normalize_identifiers(query.copy(), dialect=dialect)
+
+        # Return the hints for top level selects (model definition columns only)
+        return [
+            hint
+            for q in query.walk(prune=lambda n: not isinstance(n, exp.SetOperation))
+            if isinstance(select := q.unnest(), exp.Select)
+            for hint in _get_type_hints_for_select(
+                q, dialect, columns_to_types, start_line, end_line
+            )
+        ]
+    except Exception:
+        return []
diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py
@@ -36,6 +36,7 @@
     RenderModelRequest,
     RenderModelResponse,
 )
+from sqlmesh.lsp.hints import get_hints
 from sqlmesh.lsp.reference import get_references, get_cte_references
 from sqlmesh.lsp.uri import URI
 from web.server.api.endpoints.lineage import column_lineage, model_lineage
@@ -334,6 +335,25 @@ def hover(ls: LanguageServer, params: types.HoverParams) -> t.Optional[types.Hov
                 )
                 return None
 
+        @self.server.feature(types.TEXT_DOCUMENT_INLAY_HINT)
+        def inlay_hint(
+            ls: LanguageServer, params: types.InlayHintParams
+        ) -> t.List[types.InlayHint]:
+            """Implement type hints for sql columns as inlay hints"""
+            try:
+                uri = URI(params.text_document.uri)
+                self._ensure_context_for_document(uri)
+                if self.lsp_context is None:
+                    raise RuntimeError(f"No context found for document: {uri}")
+
+                start_line = params.range.start.line
+                end_line = params.range.end.line
+                hints = get_hints(self.lsp_context, uri, start_line, end_line)
+                return hints
+
+            except Exception as e:
+                return []
+
         @self.server.feature(types.TEXT_DOCUMENT_DEFINITION)
         def goto_definition(
             ls: LanguageServer, params: types.DefinitionParams
diff --git a/tests/lsp/test_hints.py b/tests/lsp/test_hints.py
@@ -0,0 +1,203 @@
+"""Tests for type hinting SQLMesh models"""
+
+import pytest
+
+from sqlglot import exp, parse_one
+
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.hints import get_hints, _get_type_hints_for_model_from_query
+from sqlmesh.lsp.uri import URI
+
+
+@pytest.mark.fast
+def test_hints() -> None:
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Find model URIs
+    active_customers_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.active_customers" in info.names
+    )
+    customer_revenue_lifetime_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customer_revenue_lifetime" in info.names
+    )
+    customer_revenue_by_day_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customer_revenue_by_day" in info.names
+    )
+
+    active_customers_uri = URI.from_path(active_customers_path)
+    ac_hints = get_hints(lsp_context, active_customers_uri, start_line=0, end_line=9999)
+    assert len(ac_hints) == 2
+    assert ac_hints[0].label == "::INT"
+    assert ac_hints[1].label == "::TEXT"
+
+    customer_revenue_lifetime_uri = URI.from_path(customer_revenue_lifetime_path)
+    crl_hints = get_hints(
+        lsp_context=lsp_context,
+        document_uri=customer_revenue_lifetime_uri,
+        start_line=0,
+        end_line=9999,
+    )
+    assert len(crl_hints) == 3
+    assert crl_hints[0].label == "::INT"
+    assert crl_hints[1].label == "::DOUBLE"
+    assert crl_hints[2].label == "::DATE"
+
+    customer_revenue_by_day_uri = URI.from_path(customer_revenue_by_day_path)
+    crbd_hints = get_hints(
+        lsp_context=lsp_context,
+        document_uri=customer_revenue_by_day_uri,
+        start_line=0,
+        end_line=9999,
+    )
+    assert len(crbd_hints) == 1
+    assert crbd_hints[0].label == "::INT"
+
+
+@pytest.mark.fast
+def test_union_hints() -> None:
+    query_str = """SELECT a FROM table_a UNION SELECT b FROM table_b UNION SELECT c FROM table_c"""
+    query = parse_one(query_str, dialect="postgres")
+
+    result = _get_type_hints_for_model_from_query(
+        query=query,
+        dialect="postgres",
+        columns_to_types={
+            "a": exp.DataType.build("TEXT"),
+            "b": exp.DataType.build("INT"),
+            "c": exp.DataType.build("DATE"),
+        },
+        start_line=0,
+        end_line=1,
+    )
+
+    assert len(result) == 3
+    assert result[0].label == "::DATE"
+    assert result[1].label == "::TEXT"
+    assert result[2].label == "::INT"
+
+
+@pytest.mark.fast
+def test_complex_hints() -> None:
+    query = parse_one("SELECT a, b FROM c", dialect="postgres")
+
+    result = _get_type_hints_for_model_from_query(
+        query=query,
+        dialect="postgres",
+        columns_to_types={
+            "a": exp.DataType.build("VARCHAR(100)"),
+            "b": exp.DataType.build("STRUCT<INT, STRUCT<TEXT, ARRAY<INT>>>"),
+        },
+        start_line=0,
+        end_line=1,
+    )
+
+    assert len(result) == 2
+    assert result[0].label == "::VARCHAR(100)"
+    assert result[1].label == "::STRUCT<INT, STRUCT<TEXT, INT[]>>"
+
+
+@pytest.mark.fast
+def test_simple_cast_hints() -> None:
+    """Don't add type hints if the expression is already a cast"""
+    query = parse_one("SELECT a::INT, CAST(b AS DATE), c FROM d", dialect="postgres")
+
+    result = _get_type_hints_for_model_from_query(
+        query=query,
+        dialect="postgres",
+        columns_to_types={
+            "a": exp.DataType.build("INT"),
+            "b": exp.DataType.build("DATE"),
+            "c": exp.DataType.build("TEXT"),
+        },
+        start_line=0,
+        end_line=1,
+    )
+
+    assert len(result) == 1
+    assert result[0].label == "::TEXT"
+
+
+@pytest.mark.fast
+def test_alias_cast_hints() -> None:
+    """Don't add type hints if the expression is already a cast"""
+    query = parse_one(
+        "SELECT raw_a::INT AS a, CAST(raw_b AS DATE) AS b, c FROM d", dialect="postgres"
+    )
+
+    result = _get_type_hints_for_model_from_query(
+        query=query,
+        dialect="postgres",
+        columns_to_types={
+            "a": exp.DataType.build("INT"),
+            "b": exp.DataType.build("DATE"),
+            "c": exp.DataType.build("TEXT"),
+        },
+        start_line=0,
+        end_line=1,
+    )
+
+    assert len(result) == 1
+    assert result[0].label == "::TEXT"
+
+
+@pytest.mark.fast
+def test_simple_cte_hints() -> None:
+    """Don't add type hints if the expression is already a cast"""
+    query = parse_one("WITH t AS (SELECT a FROM b) SELECT a AS c FROM t", dialect="postgres")
+
+    result = _get_type_hints_for_model_from_query(
+        query=query,
+        dialect="postgres",
+        columns_to_types={
+            "c": exp.DataType.build("INT"),
+        },
+        start_line=0,
+        end_line=1,
+    )
+
+    assert len(result) == 1
+    assert result[0].label == "::INT"
+
+
+@pytest.mark.fast
+def test_cte_with_union_hints() -> None:
+    """Don't add type hints if the expression is already a cast"""
+    query = parse_one(
+        """WITH x AS (SELECT a FROM t),
+                y AS (SELECT b FROM t),
+                z AS (SELECT c FROM t)
+         SELECT a AS d FROM x
+          UNION
+         SELECT b AS e FROM y
+          UNION
+         SELECT c AS f FROM z""",
+        dialect="postgres",
+    )
+
+    result = _get_type_hints_for_model_from_query(
+        query=query,
+        dialect="postgres",
+        columns_to_types={
+            "a": exp.DataType.build("TEXT"),
+            "b": exp.DataType.build("DATE"),
+            "c": exp.DataType.build("INT"),
+            "d": exp.DataType.build("TEXT"),
+            "e": exp.DataType.build("DATE"),
+            "f": exp.DataType.build("INT"),
+        },
+        start_line=0,
+        end_line=9999,
+    )
+
+    assert len(result) == 3
+    assert result[0].label == "::INT"
+    assert result[1].label == "::TEXT"
+    assert result[2].label == "::DATE"
diff --git a/vscode/extension/tests/hints.spec.ts b/vscode/extension/tests/hints.spec.ts
