feat(vscode): adding ability to update columns

Author: benfdking

examples/sushi/data/duckdb.db

@@ -3,13 +3,13 @@
   dialect: duckdb
   start: 1 week ago
   audits:
-    - name: not_null
-      columns: "[customer_id]"
-    - name: accepted_range
-      column: zip
-      min_v: "'00000'"
-      max_v: "'99999'"
-    - name: assert_raw_demographics
+  - name: not_null
+    columns: '[customer_id]'
+  - name: accepted_range
+    column: zip
+    min_v: "'00000'"
+    max_v: "'99999'"
+  - name: assert_raw_demographics
   columns:
     customer_id: int
     zip: text

examples/sushi/external_models.yaml

@@ -0,0 +1,42 @@
+Ability to click update columns on a model in external models
+
+# Requirements 
+
+* [ ] For each model in external models file
+* [ ] Below each model in the yaml file, provide an LSP code action to update columns
+* [ ] The code action should be able to update the columns in the model
+
+# Useful files
+
+ sqlmesh/lsp/main.py is the main file for the LSP server.
+ sqlmesh/cli/main.py and create_external_model provides similar functionality for creating external models but for the CLI and only does the whole shebang.
+
+# Implementation Plan
+
+1. **Create a new code action handler for external models YAML files**
+   - Add support for YAML files in the LSP context's `get_code_actions` method
+   - Detect when cursor is on or near a model definition in external_models.yaml
+
+2. **Add YAML file parsing to identify model positions**
+   - Parse the YAML file to identify model definitions and their line positions
+   - Track the name field and columns section for each model
+
+3. **Implement the "Update Columns" code action**
+   - When triggered, the action will:
+     - Extract the model name from the YAML
+     - Use the engine adapter to fetch current columns from the database
+     - Generate the updated columns section in YAML format
+     - Replace the existing columns section with the new one
+
+4. **Integration points**
+   - Extend `LSPContext.get_code_actions()` to handle YAML files
+   - Add YAML parsing logic to identify model boundaries
+   - Reuse existing column fetching logic from `schema_loader.py`
+   - Use the adapter's `columns()` method to get schema information
+
+5. **Code action workflow**
+   - User opens external_models.yaml in editor
+   - LSP server identifies model definitions in the file
+   - When cursor is on a model, show "Update columns" code action
+   - On activation, fetch latest schema and update the YAML
+ 

pnpm-lock.yaml

@@ -57,17 +57,6 @@ def create_external_models_file(
         external_model_fqns -= existing_model_fqns
 
     with ThreadPoolExecutor(max_workers=max_workers) as pool:
-
-        def _get_columns(table: str) -> t.Optional[t.Dict[str, t.Any]]:
-            try:
-                return adapter.columns(table, include_pseudo_columns=True)
-            except Exception as e:
-                msg = f"Unable to get schema for '{table}': '{e}'."
-                if strict:
-                    raise SQLMeshError(msg) from e
-                get_console().log_warning(msg)
-                return None
-
         gateway_part = {"gateway": gateway} if gateway else {}
 
         schemas = [
@@ -78,7 +67,7 @@ def _get_columns(table: str) -> t.Optional[t.Dict[str, t.Any]]:
             }
             for table, columns in sorted(
                 pool.map(
-                    lambda table: (table, _get_columns(table)),
+                    lambda table: (table, get_columns(adapter, dialect, table, strict)),
                     external_model_fqns,
                 )
             )
@@ -94,3 +83,20 @@ def _get_columns(table: str) -> t.Optional[t.Dict[str, t.Any]]:
 
         with open(path, "w", encoding="utf-8") as file:
             yaml.dump(entries_to_keep + schemas, file)
+
+
+def get_columns(
+    adapter: EngineAdapter, dialect: DialectType, table: str, strict: bool
+) -> t.Optional[t.Dict[str, t.Any]]:
+    """
+    Return the column and their types in a dictionary
+    """
+    try:
+        columns = adapter.columns(table, include_pseudo_columns=True)
+        return {c: dtype.sql(dialect=dialect) for c, dtype in columns.items()}
+    except Exception as e:
+        msg = f"Unable to get schema for '{table}': '{e}'."
+        if strict:
+            raise SQLMeshError(msg) from e
+        get_console().log_warning(msg)
+        return None

requirements.md

@@ -0,0 +1,7 @@
+from sqlmesh.utils.pydantic import PydanticModel
+
+EXTERNAL_MODEL_UPDATE_COLUMNS = "sqlmesh.external_model_update_columns"
+
+
+class ExternalModelUpdateColumnsRequest(PydanticModel):
+    model_name: str

sqlmesh/core/schema_loader.py

@@ -1,14 +1,22 @@
 from dataclasses import dataclass
 from pathlib import Path
+
+from pygls.server import LanguageServer
+
 from sqlmesh.core.context import Context
 import typing as t
 
-from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.core.model.definition import SqlModel, ExternalModel
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
+from sqlmesh.core.schema_loader import get_columns
+from sqlmesh.lsp.commands import EXTERNAL_MODEL_UPDATE_COLUMNS
 from sqlmesh.lsp.custom import ModelForRendering
 from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
 from sqlmesh.lsp.uri import URI
 from lsprotocol import types
+from pathlib import Path
+
+from sqlmesh.utils import yaml
 
 
 @dataclass
@@ -228,8 +236,42 @@ def get_code_actions(
                         edit=types.WorkspaceEdit(changes={params.text_document.uri: text_edits}),
                     )
                     code_actions.append(code_action)
+        return code_actions
+
+    def get_code_lenses(self, uri: URI) -> t.Optional[t.List[types.CodeLens]]:
+        models_in_file = self.map.get(uri.to_path())
+        if isinstance(models_in_file, ModelTarget):
+            models = [self.context.get_model(model) for model in models_in_file.names]
+            if any(isinstance(model, ExternalModel) for model in models):
+                code_lenses = self._get_external_model_code_lenses(uri)
+                if code_lenses:
+                    return code_lenses
+
+        return None
+
+    def _get_external_model_code_lenses(self, uri: URI) -> t.List[types.CodeLens]:
+        from sqlmesh.lsp.reference import get_yaml_model_name_ranges
+
+        """Get code actions for external models YAML files."""
+        ranges = get_yaml_model_name_ranges(uri.to_path())
+
+        code_actions: t.List[types.CodeLens] = []
+        for name, range in ranges.items():
+            # Create a code action to update columns for external models
+            command = types.Command(
+                title="Update Columns",
+                command=EXTERNAL_MODEL_UPDATE_COLUMNS,
+                arguments=[
+                    name,
+                ],
+            )
+            code_action = types.CodeLens(
+                range=range,
+                command=command,
+            )
+            code_actions.append(code_action)
 
-        return code_actions if code_actions else None
+        return code_actions if code_actions else []
 
     def list_of_models_for_rendering(self) -> t.List[ModelForRendering]:
         """Get a list of models for rendering.
@@ -332,3 +374,73 @@ def diagnostic_to_lsp_diagnostic(
             code=diagnostic.rule.name,
             code_description=types.CodeDescription(href=rule_uri),
         )
+
+    def update_external_model_columns(self, ls: LanguageServer, uri: URI, model_name: str) -> bool:
+        """
+        Update the columns for an external model in the YAML file. Returns True if changed, False if didn't because
+        of the columns already being up to date.
+
+        Errors still throw exceptions to be handled by the caller.
+        """
+        models = yaml.load(uri.to_path())
+        if not isinstance(models, list):
+            raise ValueError(
+                f"Expected a list of models in {uri.to_path()}, but got {type(models).__name__}"
+            )
+
+        existing_model = next((model for model in models if model.get("name") == model_name), None)
+        if existing_model is None:
+            raise ValueError(f"Could not find model {model_name} in {uri.to_path()}")
+
+        existing_model_columns = existing_model.get("columns")
+
+        # Get the adapter and fetch columns
+        adapter = self.context.engine_adapter
+        # Get columns for the model
+        new_columns = get_columns(
+            adapter=adapter,
+            dialect=self.context.config.model_defaults.dialect,
+            table=model_name,
+            strict=True,
+        )
+        # Compare existing columns and matching types and if they are the same, do not update
+        if existing_model_columns is not None:
+            if existing_model_columns == new_columns:
+                ls.show_message("Columns already up to date")
+                return False
+
+        # Model index to update
+        model_index = next(
+            (i for i, model in enumerate(models) if model.get("name") == model_name), None
+        )
+        if model_index is None:
+            raise ValueError(f"Could not find model {model_name} in {uri.to_path()}")
+
+        # Get end of the file to set the edit range
+        with open(uri.to_path(), "r", encoding="utf-8") as file:
+            read_file = file.read()
+
+        end_line = read_file.count("\n")
+        end_character = len(read_file.splitlines()[-1]) if end_line > 0 else 0
+
+        models[model_index]["columns"] = new_columns
+        edit = types.TextDocumentEdit(
+            text_document=types.OptionalVersionedTextDocumentIdentifier(
+                uri=uri.value,
+                version=None,
+            ),
+            edits=[
+                types.TextEdit(
+                    range=types.Range(
+                        start=types.Position(line=0, character=0),
+                        end=types.Position(
+                            line=end_line,
+                            character=end_character,
+                        ),
+                    ),
+                    new_text=yaml.dump(models),
+                )
+            ],
+        )
+        ls.apply_edit(types.WorkspaceEdit(document_changes=[edit]))
+        return True