temp

[ci skip]

Author: benfdking

examples/sushi/models/disabled.sql

@@ -3,4 +3,4 @@ MODEL (
   enabled False,
 );
 
-SELECT 1 AS a;
+SELECT 1 AS a ;

examples/sushi/models/top_waiters.sql

@@ -10,18 +10,18 @@ MODEL (
 
 WITH test_macros AS (
   SELECT
-    @ADD_ONE(1) AS  lit_two,
+    @ADD_ONE(1) AS lit_two,
     @MULTIPLY(revenue, 2.0) AS sql_exp,
     @SQL_LITERAL(revenue::TEXT, 'x', 'y', a, "b") AS sql_lit
   FROM sushi.waiter_revenue_by_day
 )
 SELECT
-  waiter_id::INT AS waiter_id ,
+  waiter_id::INT AS waiter_id,
   revenue::DOUBLE AS revenue
 FROM sushi.waiter_revenue_by_day
 WHERE
   event_date = (
-    SELECT 
+    SELECT
       MAX(event_date)
     FROM sushi.waiter_revenue_by_day
   )

sqlmesh/lsp/main.py

@@ -1,8 +1,11 @@
 #!/usr/bin/env python
 """A Language Server Protocol (LSP) server for SQL with SQLMesh integration."""
+#!/usr/bin/env python
+"""A Language Server Protocol (LSP) server for SQL with SQLMesh integration."""
 
 import asyncio
 import gc
+import io
 import logging
 import re
 import typing as t
@@ -17,7 +20,11 @@
 from lsprotocol import types
 from pygls.server import LanguageServer
 from pygls.workspace import TextDocument
+from sqlglot.errors import ParseError
 from sqlmesh.core.dialect import format_model_expressions, parse
+from sqlmesh.core.lineage import column_description
+from sqlmesh.core.model import SqlModel
+from sqlmesh.utils import type_is_known
 
 logger = logging.getLogger(__name__)
 
@@ -34,7 +41,6 @@
 """A locking mechanism for ensuring that context mutation is thread-safe."""
 
 loop = asyncio.get_event_loop()
-
 server = LanguageServer("sqlmesh-lsp", "v0.1.0", loop=loop)
 
 
@@ -93,9 +99,34 @@ async def ensure_context_for_document(document: TextDocument) -> TextDocument:
     return document
 
 
+@server.feature(types.TEXT_DOCUMENT_COMPLETION)
+async def completions(ls: LanguageServer, params: types.CompletionParams):
+    """Provide completions based on upstream model column information."""
+    items = []
+    document = await ensure_context_for_document(ls.workspace.get_document(params.text_document.uri))
+    context, model = PATHS_TO_MODELS.get(document.path, (None, None))
+    if context is None or model is None:
+        return types.CompletionList(is_incomplete=False, items=[])
+    for dep in model.depends_on:
+        model_dep = context.models[dep]
+        if model_dep.columns_to_types:
+            for column, type_ in model_dep.columns_to_types.items():
+                items.append(
+                    types.CompletionItem(
+                        label=column,
+                        label_details=types.CompletionItemLabelDetails(detail=type_.sql()),
+                        documentation=f"Source: {dep}\n\n"
+                        + (column_description(context, dep, column) or "No description available"),
+                        kind=types.CompletionItemKind.Field,
+                    )
+                )
+    return types.CompletionList(is_incomplete=False, items=items)
+
+
 @server.feature(types.TEXT_DOCUMENT_FORMATTING)
-async def formatting(ls: LanguageServer, params: types.DocumentFormattingParams):
+async def formatting(ls: LanguageServer, params: types.DocumentFormattingParams) -> t.List[types.TextEdit]:
     """Format the document based using SQLMesh format_model_expressions."""
+    logger.info(f"Formatting document: {params.text_document.uri}")
     document = await ensure_context_for_document(ls.workspace.get_document(params.text_document.uri))
     context, model = PATHS_TO_MODELS.get(document.path, (None, None))
     if context is None or model is None:
@@ -104,7 +135,8 @@ async def formatting(ls: LanguageServer, params: types.DocumentFormattingParams)
     dialect = model.dialect if model and model.is_sql else default_dialect
     try:
         expressions = parse(document.source, default_dialect=dialect)
-    except Exception:
+    except Exception as e:
+        logger.error(f"Exception occurred while parsing document: {e}")
         return []
     try:
         fmt_doc = format_model_expressions(expressions, dialect, **context.config.format.generator_options)
@@ -152,6 +184,66 @@ def _iter_match_ranges_in_projection(term: str, source: str):
             )
 
 
+def _update_diagnostics(document: TextDocument) -> None:
+    """Update diagnostics for the given document."""
+    WORKSPACE_DIAGNOSTICS[document.uri] = (document.version, diagnostics := [])
+    context, model = PATHS_TO_MODELS.get(document.path, (None, None))
+    if context is None or model is None:
+        return
+
+    default_dialect = context.default_dialect
+    dialect = model.dialect if model and model.is_sql else default_dialect
+
+    try:
+        _ = parse(document.source, default_dialect=dialect)
+    except ParseError as e:
+        for error in e.errors:
+            line = error["line"]
+            comments_before_line = [
+                _l for _l in document.lines[:line] if _l.strip().startswith(("/*", "--"))
+            ]  # This is just a hack to adjust the line number, not a proper solution but it works
+            line -= len(comments_before_line)
+            diagnostics.append(
+                types.Diagnostic(
+                    message=e.args[0],
+                    severity=types.DiagnosticSeverity.Error,
+                    range=types.Range(
+                        start=types.Position(line=line, character=error["col"]),
+                        end=types.Position(line=line, character=error["col"]),
+                    ),
+                )
+            )
+
+    if model is not None and isinstance(model, SqlModel):
+        sqlmesh_renderer_logger = logging.getLogger("sqlmesh.core.renderer")
+        buf = io.StringIO()
+        interceptor = logging.StreamHandler(stream=buf)
+        interceptor.setLevel(logging.WARNING)
+        interceptor.setFormatter(logging.Formatter("%(message)s"))
+        sqlmesh_renderer_logger.addHandler(interceptor)
+        _ = model._query_renderer.render(execution_time="now")
+        sqlmesh_renderer_logger.removeHandler(interceptor)
+        buf.seek(0)
+        warnings = buf.read().strip()
+        if warnings:
+            diagnostics.append(
+                types.Diagnostic(message=warnings, severity=types.DiagnosticSeverity.Warning, range=_top_of_file)
+            )
+
+    setattr(model, "_columns_to_types", None)  # clear cached columns to types
+    if model and model.columns_to_types:
+        for column, type_ in model.columns_to_types.items():
+            if not type_is_known(type_):
+                for range_ in _iter_match_ranges_in_projection(column, document.source):
+                    diagnostics.append(
+                        types.Diagnostic(
+                            message=f"Unknown type for column: {column} - add a type hint to final projection",
+                            severity=types.DiagnosticSeverity.Warning,
+                            range=range_,
+                        )
+                    )
+
+
 @server.feature(types.TEXT_DOCUMENT_DID_OPEN)
 async def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams):
     """Update diagnostics on document open and refresh context if necessary."""
@@ -166,6 +258,16 @@ async def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams):
                 PATHS_TO_MODELS.update(
                     {str(model._path.resolve()): (context, weakref.proxy(model)) for model in context.models.values()}
                 )
+    _update_diagnostics(document)
+    for uri, (version, diagnostics) in WORKSPACE_DIAGNOSTICS.items():
+        ls.publish_diagnostics(uri=uri, version=version, diagnostics=diagnostics)
+
+
+@server.feature(types.TEXT_DOCUMENT_DID_CLOSE)
+async def did_close(ls: LanguageServer, params: types.DidCloseTextDocumentParams):
+    """Remove diagnostics on document close."""
+    if params.text_document.uri in WORKSPACE_DIAGNOSTICS:
+        del WORKSPACE_DIAGNOSTICS[params.text_document.uri]
 
 
 @server.feature(types.TEXT_DOCUMENT_DID_SAVE)
@@ -181,6 +283,18 @@ async def did_save(ls: LanguageServer, params: types.DidOpenTextDocumentParams):
             if model._path == Path(document.path):
                 PATHS_TO_MODELS[document.path] = (context, weakref.proxy(model))
                 break
+    _update_diagnostics(document)
+    for uri, (version, diagnostics) in WORKSPACE_DIAGNOSTICS.items():
+        ls.publish_diagnostics(uri=uri, version=version, diagnostics=diagnostics)
+
+
+@server.feature(types.TEXT_DOCUMENT_DID_CHANGE)
+async def did_change(ls: LanguageServer, params: types.DidOpenTextDocumentParams):
+    """Update diagnostics on document change."""
+    document = await ensure_context_for_document(ls.workspace.get_text_document(params.text_document.uri))
+    _update_diagnostics(document)
+    for uri, (version, diagnostics) in WORKSPACE_DIAGNOSTICS.items():
+        ls.publish_diagnostics(uri=uri, version=version, diagnostics=diagnostics)
 
 
 @server.feature(types.WORKSPACE_DID_CHANGE_WATCHED_FILES)
@@ -189,6 +303,9 @@ async def did_change_watched_files(ls: LanguageServer, params: types.DidChangeWa
     updated = {}
     for change in params.changes:
         document = await ensure_context_for_document(ls.workspace.get_text_document(change.uri))
+        if change.type == types.FileChangeType.Changed:
+            _update_diagnostics(document)
+            continue
         path = Path(document.path)
         known_paths = PATHS_TO_MODELS.keys()
         if change.type == types.FileChangeType.Deleted and str(path) in known_paths:
@@ -215,8 +332,11 @@ async def did_change_watched_files(ls: LanguageServer, params: types.DidChangeWa
                     )
                     updated[context.path] = True
 
+
 def main():
+    logging.basicConfig(level=logging.DEBUG)
     server.start_io()
 
+
 if __name__ == "__main__":
     main()

vscode/extension/src/lsp/lsp.ts

@@ -1,5 +1,5 @@
 import path from "path";
-import { workspace, ExtensionContext } from "vscode";
+import { workspace, ExtensionContext, OutputChannel, window } from "vscode";
 import { ServerOptions, LanguageClientOptions, LanguageClient, TransportKind } from "vscode-languageclient/node";
 import { sqlmesh_exec, sqlmesh_lsp_exec } from "../sqlmesh/sqlmesh";
 import { err, isErr, ok, Result } from "../functional/result";
@@ -16,45 +16,42 @@ export async function activateLsp(context: ExtensionContext): Promise<Result<und
     if (workspaceFolders.length !== 1) {
         return err("Invalid number of workspace folders")
     }
+    const outputChannel: OutputChannel = window.createOutputChannel('sqlmesh_actual_lsp_implementation');
+
+    let folder = workspaceFolders[0]
     const workspacePath = workspaceFolders[0].uri.fsPath
     let serverOptions: ServerOptions = {
         run: {
             command: sqlmesh.value.bin,
             transport: TransportKind.stdio,
             options: {
-                env: {
-                    ...process.env,
-                },
                 cwd: workspacePath,
-            }
+            },
         },
         debug: {
             command: sqlmesh.value.bin,
             transport: TransportKind.stdio,
             options: {
-                env: {
-                    ...process.env,
-                },
                 cwd: workspacePath,
             }
         }
     }
     let clientOptions: LanguageClientOptions = {
         documentSelector: [
-            {
-                scheme: 'file',
-                language: 'sql',
-                pattern: '**/*.sql'
-            }
+            { scheme: 'file', pattern: `**/*.sql` }
         ],
-        synchronize: {
-            fileEvents: workspace.createFileSystemWatcher('**/*.{sql,py}'),
-        }
+        workspaceFolder: folder,
+        diagnosticCollectionName: 'sqlmesh',
+        outputChannel: outputChannel,
+        // synchronize: {
+            // fileEvents: workspace.createFileSystemWatcher('**/*.{sql,py}'),
+        // }
     }
 
-    client = new LanguageClient('sqlmesh', 'SQLMesh Language Server', serverOptions, clientOptions)
+    client = new LanguageClient('sqlmesh-lsp-example', 'SQLMesh Language Server', serverOptions, clientOptions)
     console.log('Starting language client')
-    await client.start()
+    client.start()
+
     console.log('Language client started')
     return ok(undefined)
 }