sqlmesh 0.189.1.dev2__tar.gz → 0.189.2__tar.gz

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/PKG-INFO

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

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.189.1.dev2'
-__version_tuple__ = version_tuple = (0, 189, 1, 'dev2')
+__version__ = version = '0.189.2'
+__version_tuple__ = version_tuple = (0, 189, 2)

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/core/engine_adapter/bigquery.py

@@ -199,6 +199,10 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
                 (label_tuple.expressions[0].name, label_tuple.expressions[1].name)
                 for label_tuple in label_tuples
             )
+        elif query_label_property is not None:
+            raise SQLMeshError(
+                "Invalid value for `session_properties.query_label`. Must be an array or tuple."
+            )
 
         if parsed_query_label:
             query_label_str = ",".join([":".join(label) for label in parsed_query_label])

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/core/engine_adapter/trino.py

@@ -27,6 +27,7 @@ from sqlmesh.core.engine_adapter.shared import (
     set_catalog,
 )
 from sqlmesh.core.schema_diff import SchemaDiffer
+from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.date import TimeLike
 
 if t.TYPE_CHECKING:
@@ -99,6 +100,11 @@ class TrinoEngineAdapter(
         if not isinstance(authorization, exp.Expression):
             authorization = exp.Literal.string(authorization)
 
+        if not authorization.is_string:
+            raise SQLMeshError(
+                "Invalid value for `session_properties.authorization`. Must be a string literal."
+            )
+
         authorization_sql = authorization.sql(dialect=self.dialect)
 
         self.execute(f"SET SESSION AUTHORIZATION {authorization_sql}")

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/core/linter/definition.py

@@ -8,7 +8,7 @@ import operator as op
 from collections.abc import Iterator, Iterable, Set, Mapping, Callable
 from functools import reduce
 from sqlmesh.core.model import Model
-from sqlmesh.core.linter.rule import Rule, RuleViolation
+from sqlmesh.core.linter.rule import Rule, RuleViolation, Range
 from sqlmesh.core.console import LinterConsole, get_console
 
 if t.TYPE_CHECKING:
@@ -74,6 +74,7 @@ class Linter:
                 violation_msg=violation.violation_msg,
                 model=model,
                 violation_type="error",
+                violation_range=violation.violation_range,
             )
             for violation in error_violations
         ] + [
@@ -82,6 +83,7 @@ class Linter:
                 violation_msg=violation.violation_msg,
                 model=model,
                 violation_type="warning",
+                violation_range=violation.violation_range,
             )
             for violation in warn_violations
         ]
@@ -149,7 +151,8 @@ class AnnotatedRuleViolation(RuleViolation):
         violation_msg: str,
         model: Model,
         violation_type: t.Literal["error", "warning"],
+        violation_range: t.Optional[Range] = None,
     ) -> None:
-        super().__init__(rule, violation_msg)
+        super().__init__(rule, violation_msg, violation_range)
         self.model = model
         self.violation_type = violation_type

sqlmesh-0.189.2/sqlmesh/core/linter/helpers.py

@@ -0,0 +1,115 @@
+from pathlib import Path
+
+from sqlmesh.core.linter.rule import Position, Range
+from sqlmesh.utils.pydantic import PydanticModel
+import typing as t
+
+
+class TokenPositionDetails(PydanticModel):
+    """
+    Details about a token's position in the source code in the structure provided by SQLGlot.
+
+    Attributes:
+        line (int): The line that the token ends on.
+        col (int): The column that the token ends on.
+        start (int): The start index of the token.
+        end (int): The ending index of the token.
+    """
+
+    line: int
+    col: int
+    start: int
+    end: int
+
+    @staticmethod
+    def from_meta(meta: t.Dict[str, int]) -> "TokenPositionDetails":
+        return TokenPositionDetails(
+            line=meta["line"],
+            col=meta["col"],
+            start=meta["start"],
+            end=meta["end"],
+        )
+
+    def to_range(self, read_file: t.Optional[t.List[str]]) -> Range:
+        """
+        Convert a TokenPositionDetails object to a Range object.
+
+        In the circumstances where the token's start and end positions are the same,
+        there is no need for a read_file parameter, as the range can be derived from the token's
+        line and column. This is an optimization to avoid unnecessary file reads and should
+        only be used when the token represents a single character or position in the file.
+
+        If the token's start and end positions are different, the read_file parameter is required.
+
+        :param read_file: List of lines from the file. Optional
+        :return: A Range object representing the token's position
+        """
+        if self.start == self.end:
+            # If the start and end positions are the same, we can create a range directly
+            return Range(
+                start=Position(line=self.line - 1, character=self.col - 1),
+                end=Position(line=self.line - 1, character=self.col),
+            )
+
+        if read_file is None:
+            raise ValueError("read_file must be provided when start and end positions differ.")
+
+        # Convert from 1-indexed to 0-indexed for line only
+        end_line_0 = self.line - 1
+        end_col_0 = self.col
+
+        # Find the start line and column by counting backwards from the end position
+        start_pos = self.start
+        end_pos = self.end
+
+        # Initialize with the end position
+        start_line_0 = end_line_0
+        start_col_0 = end_col_0 - (end_pos - start_pos + 1)
+
+        # If start_col_0 is negative, we need to go back to previous lines
+        while start_col_0 < 0 and start_line_0 > 0:
+            start_line_0 -= 1
+            start_col_0 += len(read_file[start_line_0])
+            # Account for newline character
+            if start_col_0 >= 0:
+                break
+            start_col_0 += 1  # For the newline character
+
+        # Ensure we don't have negative values
+        start_col_0 = max(0, start_col_0)
+        return Range(
+            start=Position(line=start_line_0, character=start_col_0),
+            end=Position(line=end_line_0, character=end_col_0),
+        )
+
+
+def read_range_from_file(file: Path, text_range: Range) -> str:
+    """
+    Read the file and return the content within the specified range.
+
+    Args:
+        file: Path to the file to read
+        text_range: The range of text to extract
+
+    Returns:
+        The content within the specified range
+    """
+    with file.open("r", encoding="utf-8") as f:
+        lines = f.readlines()
+
+    # Ensure the range is within bounds
+    start_line = max(0, text_range.start.line)
+    end_line = min(len(lines), text_range.end.line + 1)
+
+    if start_line >= end_line:
+        return ""
+
+    # Extract the relevant portions of each line
+    result = []
+    for i in range(start_line, end_line):
+        line = lines[i]
+        start_char = text_range.start.character if i == text_range.start.line else 0
+        end_char = text_range.end.character if i == text_range.end.line else len(line)
+        result.append(line[start_char:end_char])
+
+    return "".join(result)

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/core/linter/rule.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+from dataclasses import dataclass
 
 from sqlmesh.core.model import Model
 
@@ -22,6 +23,22 @@ class RuleLocation(PydanticModel):
     start_line: t.Optional[int] = None
 
 
+@dataclass(frozen=True)
+class Position:
+    """The position of a rule violation in a file, the position follows the LSP standard."""
+
+    line: int
+    character: int
+
+
+@dataclass(frozen=True)
+class Range:
+    """The range of a rule violation in a file. The range follows the LSP standard."""
+
+    start: Position
+    end: Position
+
+
 class _Rule(abc.ABCMeta):
     def __new__(cls: Type[_Rule], clsname: str, bases: t.Tuple, attrs: t.Dict) -> _Rule:
         attrs["name"] = clsname.lower()
@@ -45,9 +62,15 @@ class Rule(abc.ABC, metaclass=_Rule):
         """A summary of what this rule checks for."""
         return self.__doc__ or ""
 
-    def violation(self, violation_msg: t.Optional[str] = None) -> RuleViolation:
+    def violation(
+        self,
+        violation_msg: t.Optional[str] = None,
+        violation_range: t.Optional[Range] = None,
+    ) -> RuleViolation:
         """Create a RuleViolation instance for this rule"""
-        return RuleViolation(rule=self, violation_msg=violation_msg or self.summary)
+        return RuleViolation(
+            rule=self, violation_msg=violation_msg or self.summary, violation_range=violation_range
+        )
 
     def get_definition_location(self) -> RuleLocation:
         """Return the file path and position information for this rule.
@@ -79,9 +102,12 @@ class Rule(abc.ABC, metaclass=_Rule):
 
 
 class RuleViolation:
-    def __init__(self, rule: Rule, violation_msg: str) -> None:
+    def __init__(
+        self, rule: Rule, violation_msg: str, violation_range: t.Optional[Range] = None
+    ) -> None:
         self.rule = rule
         self.violation_msg = violation_msg
+        self.violation_range = violation_range
 
     def __repr__(self) -> str:
         return f"{self.rule.name}: {self.violation_msg}"

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/core/linter/rules/builtin.py

@@ -4,9 +4,11 @@ from __future__ import annotations
 
 import typing as t
 
+from sqlglot.expressions import Star
 from sqlglot.helper import subclasses
 
-from sqlmesh.core.linter.rule import Rule, RuleViolation
+from sqlmesh.core.linter.helpers import TokenPositionDetails
+from sqlmesh.core.linter.rule import Rule, RuleViolation, Range
 from sqlmesh.core.linter.definition import RuleSet
 from sqlmesh.core.model import Model, SqlModel
 
@@ -15,10 +17,25 @@ class NoSelectStar(Rule):
     """Query should not contain SELECT * on its outer most projections, even if it can be expanded."""
 
     def check_model(self, model: Model) -> t.Optional[RuleViolation]:
+        # Only applies to SQL models, as other model types do not have a query.
         if not isinstance(model, SqlModel):
             return None
-
-        return self.violation() if model.query.is_star else None
+        if model.query.is_star:
+            violation_range = self._get_range(model)
+            return self.violation(violation_range=violation_range)
+        return None
+
+    def _get_range(self, model: SqlModel) -> t.Optional[Range]:
+        """Get the range of the violation if available."""
+        try:
+            if len(model.query.expressions) == 1 and isinstance(model.query.expressions[0], Star):
+                return TokenPositionDetails.from_meta(model.query.expressions[0].meta).to_range(
+                    None
+                )
+        except Exception:
+            pass
+
+        return None
 
 
 class InvalidSelectStarExpansion(Rule):

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/core/model/meta.py

@@ -323,10 +323,8 @@ class ModelMeta(_Node):
 
             if prop_name == "query_label":
                 query_label = eq.right
-                if not (
-                    isinstance(query_label, exp.Array)
-                    or isinstance(query_label, exp.Tuple)
-                    or isinstance(query_label, exp.Paren)
+                if not isinstance(
+                    query_label, (exp.Array, exp.Tuple, exp.Paren, d.MacroFunc, d.MacroVar)
                 ):
                     raise ConfigError(
                         "Invalid value for `session_properties.query_label`. Must be an array or tuple."
@@ -349,7 +347,9 @@ class ModelMeta(_Node):
                         )
             elif prop_name == "authorization":
                 authorization = eq.right
-                if not (isinstance(authorization, exp.Literal) and authorization.is_string):
+                if not (
+                    isinstance(authorization, exp.Literal) and authorization.is_string
+                ) and not isinstance(authorization, (d.MacroFunc, d.MacroVar)):
                     raise ConfigError(
                         "Invalid value for `session_properties.authorization`. Must be a string literal."
                     )

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/core/plan/explainer.py

@@ -195,7 +195,7 @@ class RichExplainerConsole(ExplainerConsole):
 
                 tree.add(model_tree)
             else:
-                tree.add(f"{display_name} \[standalone audit]")
+                tree.add(f"{display_name} \\[standalone audit]")
         return tree
 
     def visit_migrate_schemas_stage(self, stage: stages.MigrateSchemasStage) -> Tree:

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/lsp/completions.py

@@ -7,14 +7,25 @@ from sqlmesh.lsp.uri import URI
 
 
 def get_sql_completions(
-    context: t.Optional[LSPContext], file_uri: t.Optional[URI]
+    context: t.Optional[LSPContext], file_uri: t.Optional[URI], content: t.Optional[str] = None
 ) -> AllModelsResponse:
     """
     Return a list of completions for a given file.
     """
+    # Get SQL keywords for the dialect
+    sql_keywords = get_keywords(context, file_uri)
+
+    # Get keywords from file content if provided
+    file_keywords = set()
+    if content:
+        file_keywords = extract_keywords_from_content(content, get_dialect(context, file_uri))
+
+    # Combine keywords - SQL keywords first, then file keywords
+    all_keywords = list(sql_keywords) + list(file_keywords - sql_keywords)
+
     return AllModelsResponse(
         models=list(get_models(context, file_uri)),
-        keywords=list(get_keywords(context, file_uri)),
+        keywords=all_keywords,
     )
 
 
@@ -97,3 +108,54 @@ def get_keywords_from_tokenizer(dialect: t.Optional[str] = None) -> t.Set[str]:
         parts = keyword.split(" ")
         expanded_keywords.update(parts)
     return expanded_keywords
+
+
+def get_dialect(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Optional[str]:
+    """
+    Get the dialect for a given file.
+    """
+    if file_uri is not None and context is not None and file_uri.to_path() in context.map:
+        file_info = context.map[file_uri.to_path()]
+
+        # Handle ModelInfo objects
+        if isinstance(file_info, ModelTarget) and file_info.names:
+            model_name = file_info.names[0]
+            model_from_context = context.context.get_model(model_name)
+            return model_from_context.dialect
+
+        # Handle AuditInfo objects
+        if isinstance(file_info, AuditTarget) and file_info.name:
+            audit = context.context.standalone_audits.get(file_info.name)
+            if audit is not None and audit.dialect:
+                return audit.dialect
+
+    if context is not None:
+        return context.context.default_dialect
+
+    return None
+
+
+def extract_keywords_from_content(content: str, dialect: t.Optional[str] = None) -> t.Set[str]:
+    """
+    Extract identifiers from SQL content using the tokenizer.
+    Only extracts identifiers (variable names, table names, column names, etc.)
+    that are not SQL keywords.
+    """
+    if not content:
+        return set()
+
+    tokenizer_class = Dialect.get_or_raise(dialect).tokenizer_class
+    keywords = set()
+    try:
+        tokenizer = tokenizer_class()
+        tokens = tokenizer.tokenize(content)
+        for token in tokens:
+            # Don't include keywords in the set
+            if token.text.upper() not in tokenizer_class.KEYWORDS:
+                keywords.add(token.text)
+
+    except Exception:
+        # If tokenization fails, return empty set
+        pass
+
+    return keywords

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/lsp/context.py

@@ -5,7 +5,7 @@ import typing as t
 
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
-from sqlmesh.lsp.custom import RenderModelEntry, ModelForRendering
+from sqlmesh.lsp.custom import ModelForRendering
 from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
 from sqlmesh.lsp.uri import URI
 
@@ -176,15 +176,18 @@ class LSPContext:
             if audit._path is not None
         ]
 
-    def get_autocomplete(self, uri: t.Optional[URI]) -> AllModelsResponse:
+    def get_autocomplete(
+        self, uri: t.Optional[URI], content: t.Optional[str] = None
+    ) -> AllModelsResponse:
         """Get autocomplete suggestions for a file.
 
         Args:
             uri: The URI of the file to get autocomplete suggestions for.
+            content: The content of the file (optional).
 
         Returns:
             AllModelsResponse containing models and keywords.
         """
         from sqlmesh.lsp.completions import get_sql_completions
 
-        return get_sql_completions(self, uri)
+        return get_sql_completions(self, uri, content)

{sqlmesh-0.189.1.dev2 → sqlmesh-0.189.2}/sqlmesh/lsp/main.py

@@ -36,9 +36,7 @@ from sqlmesh.lsp.custom import (
     RenderModelRequest,
     RenderModelResponse,
 )
-from sqlmesh.lsp.reference import (
-    get_references,
-)
+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
 from web.server.api.endpoints.models import get_models
@@ -114,13 +112,22 @@ class SQLMeshLanguageServer:
         @self.server.feature(ALL_MODELS_FEATURE)
         def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsResponse:
             uri = URI(params.textDocument.uri)
+
+            # Get the document content
+            content = None
+            try:
+                document = ls.workspace.get_text_document(params.textDocument.uri)
+                content = document.source
+            except Exception:
+                pass
+
             try:
                 context = self._context_get_or_load(uri)
-                return context.get_autocomplete(uri)
+                return context.get_autocomplete(uri, content)
             except Exception as e:
                 from sqlmesh.lsp.completions import get_sql_completions
 
-                return get_sql_completions(None, URI(params.textDocument.uri))
+                return get_sql_completions(None, URI(params.textDocument.uri), content)
 
         @self.server.feature(RENDER_MODEL_FEATURE)
         def render_model(ls: LanguageServer, params: RenderModelRequest) -> RenderModelResponse:
@@ -369,6 +376,28 @@ class SQLMeshLanguageServer:
                 ls.show_message(f"Error getting references: {e}", types.MessageType.Error)
                 return []
 
+        @self.server.feature(types.TEXT_DOCUMENT_REFERENCES)
+        def find_references(
+            ls: LanguageServer, params: types.ReferenceParams
+        ) -> t.Optional[t.List[types.Location]]:
+            """Find all references of a symbol (currently supporting CTEs)"""
+            try:
+                uri = URI(params.text_document.uri)
+                self._ensure_context_for_document(uri)
+                document = ls.workspace.get_text_document(params.text_document.uri)
+                if self.lsp_context is None:
+                    raise RuntimeError(f"No context found for document: {document.path}")
+
+                cte_references = get_cte_references(self.lsp_context, uri, params.position)
+
+                # Convert references to Location objects
+                locations = [types.Location(uri=ref.uri, range=ref.range) for ref in cte_references]
+
+                return locations if locations else None
+            except Exception as e:
+                ls.show_message(f"Error getting locations: {e}", types.MessageType.Error)
+                return None
+
         @self.server.feature(types.TEXT_DOCUMENT_DIAGNOSTIC)
         def diagnostic(
             ls: LanguageServer, params: types.DocumentDiagnosticParams
@@ -471,8 +500,16 @@ class SQLMeshLanguageServer:
                 uri = URI(params.text_document.uri)
                 context = self._context_get_or_load(uri)
 
+                # Get the document content
+                content = None
+                try:
+                    document = ls.workspace.get_text_document(params.text_document.uri)
+                    content = document.source
+                except Exception:
+                    pass
+
                 # Get completions using the existing completions module
-                completion_response = context.get_autocomplete(uri)
+                completion_response = context.get_autocomplete(uri, content)
 
                 completion_items = []
                 # Add model completions
@@ -618,8 +655,24 @@ class SQLMeshLanguageServer:
     ) -> t.Optional[types.Diagnostic]:
         if diagnostic.model._path is None:
             return None
-        with open(diagnostic.model._path, "r", encoding="utf-8") as file:
-            lines = file.readlines()
+        if not diagnostic.violation_range:
+            with open(diagnostic.model._path, "r", encoding="utf-8") as file:
+                lines = file.readlines()
+            range = types.Range(
+                start=types.Position(line=0, character=0),
+                end=types.Position(line=len(lines) - 1, character=len(lines[-1])),
+            )
+        else:
+            range = types.Range(
+                start=types.Position(
+                    line=diagnostic.violation_range.start.line,
+                    character=diagnostic.violation_range.start.character,
+                ),
+                end=types.Position(
+                    line=diagnostic.violation_range.end.line,
+                    character=diagnostic.violation_range.end.character,
+                ),
+            )
 
         # Get rule definition location for diagnostics link
         rule_location = diagnostic.rule.get_definition_location()
@@ -628,10 +681,7 @@ class SQLMeshLanguageServer:
 
         # Use URI format to create a link for "related information"
         return types.Diagnostic(
-            range=types.Range(
-                start=types.Position(line=0, character=0),
-                end=types.Position(line=len(lines), character=len(lines[-1])),
-            ),
+            range=range,
             message=diagnostic.violation_msg,
             severity=types.DiagnosticSeverity.Error
             if diagnostic.violation_type == "error"