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

@@ -4,38 +4,33 @@
 import asyncio
 import gc
 import logging
-import re
 import typing as t
 import weakref
 from collections import defaultdict
 from contextlib import suppress
-from functools import lru_cache
 from itertools import cycle
 from pathlib import Path
 
 import sqlmesh
 from lsprotocol import types
 from pygls.server import LanguageServer
 from pygls.workspace import TextDocument
-from sqlmesh.core.dialect import format_model_expressions, parse
+from sqlmesh._version import __version__
 
 logger = logging.getLogger(__name__)
 
-WORKSPACE_DIAGNOSTICS: t.Dict[str, t.Tuple[t.Optional[int], t.List[types.Diagnostic]]] = {}
-"""A mapping of document URIs to diagnostics."""
-
 CONTEXTS: t.Dict[str, sqlmesh.Context] = {}
 """A mapping of workspace paths to SQLMesh contexts."""
 
 PATHS_TO_MODELS: t.Dict[str, t.Tuple[sqlmesh.Context, sqlmesh.Model]] = {}
 """A mapping of file paths to SQLMesh (context, model) tuples."""
 
-C_MUTEX = defaultdict(asyncio.Lock)
+C_MUTEX: t.DefaultDict[t.Union[str, Path], asyncio.Lock] = defaultdict(asyncio.Lock)
 """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)
+server = LanguageServer("sqlmesh_lsp", __version__, loop=loop)
 
 
 async def refresh_context_loop(context: sqlmesh.Context) -> None:
@@ -46,11 +41,15 @@ async def refresh_context_loop(context: sqlmesh.Context) -> None:
     gc_iter = cycle(list(range(10)))
     while True:
         await asyncio.sleep(10.0)
-        if context._loader.reload_needed():
-            async with C_MUTEX[context.path]:
-                await asyncio.to_thread(context.load)
+        for loader in context._loaders:
+            if loader.reload_needed():
+                async with C_MUTEX[context.path]:
+                    await asyncio.to_thread(context.load)
                 PATHS_TO_MODELS.update(
-                    {str(model._path.resolve()): (context, weakref.proxy(model)) for model in context.models.values()}
+                    {
+                        str(model._path.resolve()): (context, weakref.proxy(model))
+                        for model in context.models.values()
+                    }
                 )
         if next(gc_iter) == 0:
             gc.collect()
@@ -84,7 +83,10 @@ async def ensure_context_for_document(document: TextDocument) -> TextDocument:
                     loop.create_task(refresh_context_loop(handle))
                     CONTEXTS[str(path)] = handle
                     PATHS_TO_MODELS.update(
-                        {str(model._path.resolve()): (handle, weakref.proxy(model)) for model in handle.models.values()}
+                        {
+                            str(model._path.resolve()): (handle, weakref.proxy(model))
+                            for model in handle.models.values()
+                        }
                     )
                     server.show_message(f"Context loaded for: {path}")
                     loaded = True
@@ -94,99 +96,81 @@ async def ensure_context_for_document(document: TextDocument) -> TextDocument:
 
 
 @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."""
-    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 []
-    default_dialect = context.default_dialect
-    dialect = model.dialect if model and model.is_sql else default_dialect
-    try:
-        expressions = parse(document.source, default_dialect=dialect)
-    except Exception:
-        return []
     try:
-        fmt_doc = format_model_expressions(expressions, dialect, **context.config.format.generator_options)
-        if context.config.format.append_newline:
-            fmt_doc += "\n"
+        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))
+        context.format(paths=[Path(document.path)])
+        with open(document.path, "r+", encoding="utf-8") as file:
+            return [
+                types.TextEdit(
+                    range=types.Range(
+                        types.Position(0, 0),
+                        types.Position(len(document.lines), len(document.lines[-1])),
+                    ),
+                    new_text=file.read(),
+                )
+            ]
     except Exception as e:
         ls.show_message(f"Error formatting SQL: {e}", types.MessageType.Error)
         return []
-    return [
-        types.TextEdit(
-            range=types.Range(
-                types.Position(0, 0),
-                types.Position(len(document.lines), len(document.lines[-1])),
-            ),
-            new_text=fmt_doc,
-        )
-    ]
-
-
-_top_of_file = types.Range(start=types.Position(line=0, character=0), end=types.Position(line=0, character=0))
-
-_cached_re_compile = t.cast(t.Callable[[str, re.RegexFlag], re.Pattern[str]], lru_cache(maxsize=1024)(re.compile))
-
-
-def _iter_match_ranges_in_projection(term: str, source: str):
-    """Iterate over ranges of matches for a term in a SQL projection."""
-    col_patt = _cached_re_compile(rf'\b["`]?({term})["`]?,?', re.IGNORECASE)
-    projection_patt = _cached_re_compile(r"SELECT\s+(.*)\s+FROM", re.DOTALL | re.IGNORECASE)
-    for p_match in projection_patt.finditer(source):
-        if not p_match.group(1):
-            continue
-        proj_start, proj_end = p_match.span(1)
-        proj_substr = source[proj_start:proj_end]
-        for c_match in col_patt.finditer(proj_substr):
-            if not c_match.group(1):
-                continue
-            col_start, col_end = c_match.span(1)
-            start, end = proj_start + col_start, proj_start + col_end
-            line = source.count("\n", 0, start)
-            char_s = start - source.rfind("\n", 0, start) - 1
-            char_e = end - source.rfind("\n", 0, end)
-            yield types.Range(
-                start=types.Position(line=line, character=char_s),
-                end=types.Position(line=line, character=char_e),
-            )
 
 
 @server.feature(types.TEXT_DOCUMENT_DID_OPEN)
-async def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams):
+async def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> None:
     """Update diagnostics on document open and refresh context if necessary."""
-    document = await ensure_context_for_document(ls.workspace.get_document(params.text_document.uri))
+    document = await ensure_context_for_document(
+        ls.workspace.get_document(params.text_document.uri)
+    )
     path = Path(document.path)
     known_paths = PATHS_TO_MODELS.keys()
     for context in CONTEXTS.values():
-        if path.is_relative_to(context.path) and path.suffix in (".sql", ".py") and str(path) not in known_paths:
+        if (
+            path.is_relative_to(context.path)
+            and path.suffix in (".sql", ".py")
+            and str(path) not in known_paths
+        ):
             ls.show_message(f"Refreshing context with new file: {path}", types.MessageType.Info)
             async with C_MUTEX[context.path]:
                 await asyncio.to_thread(context.load)
                 PATHS_TO_MODELS.update(
-                    {str(model._path.resolve()): (context, weakref.proxy(model)) for model in context.models.values()}
+                    {
+                        str(model._path.resolve()): (context, weakref.proxy(model))
+                        for model in context.models.values()
+                    }
                 )
 
 
 @server.feature(types.TEXT_DOCUMENT_DID_SAVE)
-async def did_save(ls: LanguageServer, params: types.DidOpenTextDocumentParams):
+async def did_save(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> None:
     """Update diagnostics on document save."""
-    document = await ensure_context_for_document(ls.workspace.get_document(params.text_document.uri))
+    document = await ensure_context_for_document(
+        ls.workspace.get_document(params.text_document.uri)
+    )
     context, _ = PATHS_TO_MODELS.get(document.path, (None, None))
     if context is not None:
-        context._loader._path_mtimes[Path(document.path)] = 0.0
-        async with C_MUTEX[context.path]:
-            await asyncio.to_thread(context.load)
-        for model in context.models.values():
-            if model._path == Path(document.path):
-                PATHS_TO_MODELS[document.path] = (context, weakref.proxy(model))
-                break
+        for loader in context._loaders:
+            loader._path_mtimes[Path(document.path)] = 0.0
+            async with C_MUTEX[context.path]:
+                await asyncio.to_thread(context.load)
+            for model in context.models.values():
+                if model._path == Path(document.path):
+                    PATHS_TO_MODELS[document.path] = (context, weakref.proxy(model))
+                    break
 
 
 @server.feature(types.WORKSPACE_DID_CHANGE_WATCHED_FILES)
-async def did_change_watched_files(ls: LanguageServer, params: types.DidChangeWatchedFilesParams):
+async def did_change_watched_files(
+    ls: LanguageServer, params: types.DidChangeWatchedFilesParams
+) -> None:
     """Refresh context if a file changes."""
-    updated = {}
+    updated: t.Dict[t.Union[str, Path], bool] = {}
     for change in params.changes:
         document = await ensure_context_for_document(ls.workspace.get_text_document(change.uri))
         path = Path(document.path)
@@ -215,8 +199,11 @@ async def did_change_watched_files(ls: LanguageServer, params: types.DidChangeWa
                     )
                     updated[context.path] = True
 
-def main():
+
+def main() -> None:
+    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)
 }