sqlmesh 0.187.6.dev5__py3-none-any.whl → 0.187.6.dev10__py3-none-any.whl

sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.187.6.dev5'
-__version_tuple__ = version_tuple = (0, 187, 6, 'dev5')
+__version__ = version = '0.187.6.dev10'
+__version_tuple__ = version_tuple = (0, 187, 6, 'dev10')

sqlmesh/core/context.py

@@ -932,19 +932,22 @@ class GenericContext(BaseContext, t.Generic[C]):
 
         return snapshot
 
-    def config_for_path(self, path: Path) -> Config:
+    def config_for_path(self, path: Path) -> t.Tuple[Config, Path]:
+        """Returns the config and path of the said project for a given file path."""
         for config_path, config in self.configs.items():
             try:
                 path.relative_to(config_path)
-                return config
+                return config, config_path
             except ValueError:
                 pass
-        return self.config
+        return self.config, self.path
 
     def config_for_node(self, node: str | Model | Audit) -> Config:
         if isinstance(node, str):
-            return self.config_for_path(self.get_snapshot(node, raise_if_missing=True).node._path)  # type: ignore
-        return self.config_for_path(node._path)  # type: ignore
+            return self.config_for_path(self.get_snapshot(node, raise_if_missing=True).node._path)[
+                0
+            ]  # type: ignore
+        return self.config_for_path(node._path)[0]  # type: ignore
 
     @property
     def models(self) -> MappingProxyType[str, Model]:

sqlmesh/core/selector.py

@@ -93,7 +93,7 @@ class Selector:
             }
 
         all_selected_models = self.expand_model_selections(
-            model_selections, models={**self._models, **env_models}
+            model_selections, models={**env_models, **self._models}
         )
 
         dag: DAG[str] = DAG()

sqlmesh/lsp/main.py

@@ -57,8 +57,12 @@ class SQLMeshLanguageServer:
         self.server = LanguageServer(server_name, version)
         self.context_class = context_class
         self.lsp_context: t.Optional[LSPContext] = None
-        self.lint_cache: t.Dict[URI, t.List[AnnotatedRuleViolation]] = {}
 
+        # Cache stores tuples of (diagnostics, diagnostic_version)
+        self.lint_cache: t.Dict[URI, t.Tuple[t.List[AnnotatedRuleViolation], int]] = {}
+        self._diagnostic_version_counter: int = 0
+
+        self.client_supports_pull_diagnostics = False
         # Register LSP features (e.g., formatting, hover, etc.)
         self._register_features()
 
@@ -69,6 +73,18 @@ class SQLMeshLanguageServer:
         def initialize(ls: LanguageServer, params: types.InitializeParams) -> None:
             """Initialize the server when the client connects."""
             try:
+                # Check if client supports pull diagnostics
+                if params.capabilities and params.capabilities.text_document:
+                    diagnostics = getattr(params.capabilities.text_document, "diagnostic", None)
+                    if diagnostics:
+                        self.client_supports_pull_diagnostics = True
+                        ls.log_trace("Client supports pull diagnostics")
+                    else:
+                        self.client_supports_pull_diagnostics = False
+                        ls.log_trace("Client does not support pull diagnostics")
+                else:
+                    self.client_supports_pull_diagnostics = False
+
                 if params.workspace_folders:
                     # Try to find a SQLMesh config file in any workspace folder (only at the root level)
                     for folder in params.workspace_folders:
@@ -153,61 +169,71 @@ class SQLMeshLanguageServer:
         def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> None:
             uri = URI(params.text_document.uri)
             context = self._context_get_or_load(uri)
-            if self.lint_cache.get(uri) is not None:
+            models = context.map[uri.to_path()]
+            if models is None or not isinstance(models, ModelTarget):
+                return
+
+            if self.lint_cache.get(uri) is None:
+                diagnostics = context.context.lint_models(
+                    models.names,
+                    raise_on_error=False,
+                )
+                self._diagnostic_version_counter += 1
+                self.lint_cache[uri] = (diagnostics, self._diagnostic_version_counter)
+
+            # Only publish diagnostics if client doesn't support pull diagnostics
+            if not self.client_supports_pull_diagnostics:
+                diagnostics, _ = self.lint_cache[uri]
                 ls.publish_diagnostics(
                     params.text_document.uri,
-                    SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(self.lint_cache[uri]),
+                    SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
                 )
-                return
-            models = context.map[uri.to_path()]
-            if models is None:
-                return
-            if not isinstance(models, ModelTarget):
-                return
-            self.lint_cache[uri] = context.context.lint_models(
-                models.names,
-                raise_on_error=False,
-            )
-            ls.publish_diagnostics(
-                params.text_document.uri,
-                SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(self.lint_cache[uri]),
-            )
 
         @self.server.feature(types.TEXT_DOCUMENT_DID_CHANGE)
         def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) -> None:
             uri = URI(params.text_document.uri)
             context = self._context_get_or_load(uri)
             models = context.map[uri.to_path()]
-            if models is None:
+            if models is None or not isinstance(models, ModelTarget):
                 return
-            if not isinstance(models, ModelTarget):
-                return
-            self.lint_cache[uri] = context.context.lint_models(
+
+            # Always update the cache
+            diagnostics = context.context.lint_models(
                 models.names,
                 raise_on_error=False,
             )
-            ls.publish_diagnostics(
-                params.text_document.uri,
-                SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(self.lint_cache[uri]),
-            )
+            self._diagnostic_version_counter += 1
+            self.lint_cache[uri] = (diagnostics, self._diagnostic_version_counter)
+
+            # Only publish diagnostics if client doesn't support pull diagnostics
+            if not self.client_supports_pull_diagnostics:
+                ls.publish_diagnostics(
+                    params.text_document.uri,
+                    SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
+                )
 
         @self.server.feature(types.TEXT_DOCUMENT_DID_SAVE)
         def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> None:
             uri = URI(params.text_document.uri)
             context = self._context_get_or_load(uri)
             models = context.map[uri.to_path()]
-            if models is None:
-                return
-            if not isinstance(models, ModelTarget):
+            if models is None or not isinstance(models, ModelTarget):
                 return
-            self.lint_cache[uri] = context.context.lint_models(
+
+            # Always update the cache
+            diagnostics = context.context.lint_models(
                 models.names,
                 raise_on_error=False,
             )
-            ls.publish_diagnostics(
-                params.text_document.uri,
-                SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(self.lint_cache[uri]),
-            )
+            self._diagnostic_version_counter += 1
+            self.lint_cache[uri] = (diagnostics, self._diagnostic_version_counter)
+
+            # Only publish diagnostics if client doesn't support pull diagnostics
+            if not self.client_supports_pull_diagnostics:
+                ls.publish_diagnostics(
+                    params.text_document.uri,
+                    SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
+                )
 
         @self.server.feature(types.TEXT_DOCUMENT_FORMATTING)
         def formatting(
@@ -327,6 +353,125 @@ class SQLMeshLanguageServer:
                 ls.show_message(f"Error getting references: {e}", types.MessageType.Error)
                 return []
 
+        @self.server.feature(types.TEXT_DOCUMENT_DIAGNOSTIC)
+        def diagnostic(
+            ls: LanguageServer, params: types.DocumentDiagnosticParams
+        ) -> types.DocumentDiagnosticReport:
+            """Handle diagnostic pull requests from the client."""
+            try:
+                uri = URI(params.text_document.uri)
+                diagnostics, result_id = self._get_diagnostics_for_uri(uri)
+
+                # Check if client provided a previous result ID
+                if hasattr(params, "previous_result_id") and params.previous_result_id == result_id:
+                    # Return unchanged report if diagnostics haven't changed
+                    return types.RelatedUnchangedDocumentDiagnosticReport(
+                        kind=types.DocumentDiagnosticReportKind.Unchanged,
+                        result_id=str(result_id),
+                    )
+
+                return types.RelatedFullDocumentDiagnosticReport(
+                    kind=types.DocumentDiagnosticReportKind.Full,
+                    items=diagnostics,
+                    result_id=str(result_id),
+                )
+            except Exception as e:
+                ls.show_message(f"Error getting diagnostics: {e}", types.MessageType.Error)
+                return types.RelatedFullDocumentDiagnosticReport(
+                    kind=types.DocumentDiagnosticReportKind.Full,
+                    items=[],
+                )
+
+        @self.server.feature(types.WORKSPACE_DIAGNOSTIC)
+        def workspace_diagnostic(
+            ls: LanguageServer, params: types.WorkspaceDiagnosticParams
+        ) -> types.WorkspaceDiagnosticReport:
+            """Handle workspace-wide diagnostic pull requests from the client."""
+            try:
+                if self.lsp_context is None:
+                    current_path = Path.cwd()
+                    self._ensure_context_in_folder(current_path)
+
+                if self.lsp_context is None:
+                    return types.WorkspaceDiagnosticReport(items=[])
+
+                items: t.List[
+                    t.Union[
+                        types.WorkspaceFullDocumentDiagnosticReport,
+                        types.WorkspaceUnchangedDocumentDiagnosticReport,
+                    ]
+                ] = []
+
+                # Get all SQL and Python model files from the context
+                for path, target in self.lsp_context.map.items():
+                    if isinstance(target, ModelTarget):
+                        uri = URI.from_path(path)
+                        diagnostics, result_id = self._get_diagnostics_for_uri(uri)
+
+                        # Check if we have a previous result ID for this file
+                        previous_result_id = None
+                        if hasattr(params, "previous_result_ids") and params.previous_result_ids:
+                            for prev in params.previous_result_ids:
+                                if prev.uri == uri.value:
+                                    previous_result_id = prev.value
+                                    break
+
+                        if previous_result_id and previous_result_id == result_id:
+                            # File hasn't changed
+                            items.append(
+                                types.WorkspaceUnchangedDocumentDiagnosticReport(
+                                    kind=types.DocumentDiagnosticReportKind.Unchanged,
+                                    result_id=str(result_id),
+                                    uri=uri.value,
+                                )
+                            )
+                        else:
+                            # File has changed or is new
+                            items.append(
+                                types.WorkspaceFullDocumentDiagnosticReport(
+                                    kind=types.DocumentDiagnosticReportKind.Full,
+                                    result_id=str(result_id),
+                                    uri=uri.value,
+                                    items=diagnostics,
+                                )
+                            )
+
+                return types.WorkspaceDiagnosticReport(items=items)
+
+            except Exception as e:
+                ls.show_message(
+                    f"Error getting workspace diagnostics: {e}", types.MessageType.Error
+                )
+                return types.WorkspaceDiagnosticReport(items=[])
+
+    def _get_diagnostics_for_uri(self, uri: URI) -> t.Tuple[t.List[types.Diagnostic], int]:
+        """Get diagnostics for a specific URI, returning (diagnostics, result_id)."""
+        # Check if we have cached diagnostics
+        if uri in self.lint_cache:
+            diagnostics, result_id = self.lint_cache[uri]
+            return SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics), result_id
+
+        # Try to get diagnostics by loading context and linting
+        try:
+            context = self._context_get_or_load(uri)
+            models = context.map[uri.to_path()]
+            if models is None or not isinstance(models, ModelTarget):
+                return [], 0
+
+            # Lint the models and cache the results
+            diagnostics = context.context.lint_models(
+                models.names,
+                raise_on_error=False,
+            )
+            self._diagnostic_version_counter += 1
+            self.lint_cache[uri] = (diagnostics, self._diagnostic_version_counter)
+            return SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(
+                diagnostics
+            ), self._diagnostic_version_counter
+        except Exception:
+            # If we can't get diagnostics, return empty list with no result ID
+            return [], 0
+
     def _context_get_or_load(self, document_uri: URI) -> LSPContext:
         if self.lsp_context is None:
             self._ensure_context_for_document(document_uri)

sqlmesh/lsp/reference.py

@@ -1,6 +1,8 @@
 from lsprotocol.types import Range, Position
 import typing as t
+from pathlib import Path
 
+from sqlmesh.core.audit import StandaloneAudit
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
@@ -10,6 +12,10 @@ from sqlglot.optimizer.scope import build_scope
 from sqlmesh.lsp.uri import URI
 from sqlmesh.utils.pydantic import PydanticModel
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+import ast
+from sqlmesh.core.model import Model
+from sqlmesh import macro
+import inspect
 
 
 class Reference(PydanticModel):
@@ -72,6 +78,11 @@ def get_references(
         A list of references at the given position
     """
     references = get_model_definitions_for_a_path(lint_context, document_uri)
+
+    # Get macro references before filtering by position
+    macro_references = get_macro_definitions_for_a_path(lint_context, document_uri)
+    references.extend(macro_references)
+
     filtered_references = list(filter(by_position(position), references))
     return filtered_references
 
@@ -134,9 +145,12 @@ def get_model_definitions_for_a_path(
     with open(file_path, "r", encoding="utf-8") as file:
         read_file = file.readlines()
 
-    # Build scope tree to properly handle nested CTEs
-    query = normalize_identifiers(query.copy(), dialect=dialect)
-    root_scope = build_scope(query)
+    # Build a scope tree to properly handle nested CTEs
+    try:
+        query = normalize_identifiers(query.copy(), dialect=dialect)
+        root_scope = build_scope(query)
+    except Exception:
+        root_scope = None
 
     if root_scope:
         # Traverse all scopes to find CTE definitions and table references
@@ -287,3 +301,180 @@ def _range_from_token_position_details(
         start=Position(line=start_line_0, character=start_col_0),
         end=Position(line=end_line_0, character=end_col_0),
     )
+
+
+def get_macro_definitions_for_a_path(
+    lsp_context: LSPContext, document_uri: URI
+) -> t.List[Reference]:
+    """
+    Get macro references for a given path.
+
+    This function finds all macro invocations (e.g., @ADD_ONE, @MULTIPLY) in a SQL file
+    and creates references to their definitions in the Python macro files.
+
+    Args:
+        lsp_context: The LSP context containing macro definitions
+        document_uri: The URI of the document to search for macro invocations
+
+    Returns:
+        A list of Reference objects for each macro invocation found
+    """
+    path = document_uri.to_path()
+    if path.suffix != ".sql":
+        return []
+
+    # Get the file info from the context map
+    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 isinstance(file_info, ModelTarget):
+        # It's a model
+        target: t.Optional[t.Union[Model, StandaloneAudit]] = lsp_context.context.get_model(
+            model_or_snapshot=file_info.names[0], raise_if_missing=False
+        )
+        if target is None or not isinstance(target, SqlModel):
+            return []
+        query = target.query
+        file_path = target._path
+    elif isinstance(file_info, AuditTarget):
+        # It's a standalone audit
+        target = lsp_context.context.standalone_audits.get(file_info.name)
+        if target is None:
+            return []
+        query = target.query
+        file_path = target._path
+    else:
+        return []
+
+    references = []
+    config_for_model, config_path = lsp_context.context.config_for_path(
+        file_path,
+    )
+
+    with open(file_path, "r", encoding="utf-8") as file:
+        read_file = file.readlines()
+
+    for node in query.find_all(exp.Anonymous):
+        macro_name = node.name.lower()
+        reference = get_macro_reference(
+            node=node,
+            target=target,
+            read_file=read_file,
+            config_path=config_path,
+            macro_name=macro_name,
+        )
+        if reference is not None:
+            references.append(reference)
+
+    return references
+
+
+def get_macro_reference(
+    target: t.Union[Model, StandaloneAudit],
+    read_file: t.List[str],
+    config_path: t.Optional[Path],
+    node: exp.Expression,
+    macro_name: str,
+) -> t.Optional[Reference]:
+    # Get the file path where the macro is defined
+    try:
+        # Get the position of the macro invocation in the source file first
+        if hasattr(node, "meta") and node.meta:
+            token_details = TokenPositionDetails.from_meta(node.meta)
+            macro_range = _range_from_token_position_details(token_details, read_file)
+
+            # Check if it's a built-in method
+            if builtin := get_built_in_macro_reference(macro_name, macro_range):
+                return builtin
+        else:
+            # Skip if we can't get the position
+            return None
+
+        # Find the macro definition information
+        macro_def = target.python_env.get(macro_name)
+        if macro_def is None:
+            return None
+
+        function_name = macro_def.name
+        if not function_name:
+            return None
+        if not macro_def.path:
+            return None
+        if not config_path:
+            return None
+        path = Path(config_path).joinpath(macro_def.path)
+
+        # Parse the Python file to find the function definition
+        with open(path, "r") as f:
+            tree = ast.parse(f.read())
+        with open(path, "r") as f:
+            output_read_line = f.readlines()
+
+        # Find the function definition by name
+        start_line = None
+        end_line = None
+        get_length_of_end_line = None
+        docstring = None
+        for ast_node in ast.walk(tree):
+            if isinstance(ast_node, ast.FunctionDef) and ast_node.name == function_name:
+                start_line = ast_node.lineno
+                end_line = ast_node.end_lineno
+                get_length_of_end_line = (
+                    len(output_read_line[end_line - 1])
+                    if end_line is not None and end_line - 1 < len(read_file)
+                    else 0
+                )
+                # Extract docstring if present
+                docstring = ast.get_docstring(ast_node)
+                break
+
+        if start_line is None or end_line is None or get_length_of_end_line is None:
+            return None
+
+        # Create a reference to the macro definition
+        macro_uri = URI.from_path(path)
+
+        return Reference(
+            uri=macro_uri.value,
+            range=macro_range,
+            target_range=Range(
+                start=Position(line=start_line - 1, character=0),
+                end=Position(line=end_line - 1, character=get_length_of_end_line),
+            ),
+            markdown_description=docstring,
+        )
+    except Exception:
+        return None
+
+
+def get_built_in_macro_reference(macro_name: str, macro_range: Range) -> t.Optional[Reference]:
+    """
+    Get a reference to a built-in macro by its name.
+
+    Args:
+        macro_name: The name of the built-in macro (e.g., 'each', 'sql_literal')
+        macro_range: The range of the macro invocation in the source file
+    """
+    built_in_macros = macro.get_registry()
+    built_in_macro = built_in_macros.get(macro_name)
+    if built_in_macro is None:
+        return None
+
+    func = built_in_macro.func
+    filename = inspect.getfile(func)
+    source_lines, line_number = inspect.getsourcelines(func)
+
+    # Calculate the end line number by counting the number of source lines
+    end_line_number = line_number + len(source_lines) - 1
+
+    return Reference(
+        uri=URI.from_path(Path(filename)).value,
+        range=macro_range,
+        target_range=Range(
+            start=Position(line=line_number - 1, character=0),
+            end=Position(line=end_line_number - 1, character=0),
+        ),
+        markdown_description=func.__doc__ if func.__doc__ else None,
+    )

{sqlmesh-0.187.6.dev5.dist-info → sqlmesh-0.187.6.dev10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.187.6.dev5
+Version: 0.187.6.dev10
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License