sqlmesh 0.196.4.dev2__tar.gz → 0.196.4.dev3__tar.gz

{sqlmesh-0.196.4.dev2 → sqlmesh-0.196.4.dev3}/PKG-INFO

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

{sqlmesh-0.196.4.dev2 → sqlmesh-0.196.4.dev3}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.196.4.dev2'
-__version_tuple__ = version_tuple = (0, 196, 4, 'dev2')
+__version__ = version = '0.196.4.dev3'
+__version_tuple__ = version_tuple = (0, 196, 4, 'dev3')

{sqlmesh-0.196.4.dev2 → sqlmesh-0.196.4.dev3}/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, Range
+from sqlmesh.core.linter.rule import Rule, RuleViolation, Range, Fix
 from sqlmesh.core.console import LinterConsole, get_console
 
 if t.TYPE_CHECKING:
@@ -75,6 +75,7 @@ class Linter:
                 model=model,
                 violation_type="error",
                 violation_range=violation.violation_range,
+                fixes=violation.fixes,
             )
             for violation in error_violations
         ] + [
@@ -84,6 +85,7 @@ class Linter:
                 model=model,
                 violation_type="warning",
                 violation_range=violation.violation_range,
+                fixes=violation.fixes,
             )
             for violation in warn_violations
         ]
@@ -152,7 +154,8 @@ class AnnotatedRuleViolation(RuleViolation):
         model: Model,
         violation_type: t.Literal["error", "warning"],
         violation_range: t.Optional[Range] = None,
+        fixes: t.Optional[t.List[Fix]] = None,
     ) -> None:
-        super().__init__(rule, violation_msg, violation_range)
+        super().__init__(rule, violation_msg, violation_range, fixes)
         self.model = model
         self.violation_type = violation_type

{sqlmesh-0.196.4.dev2 → sqlmesh-0.196.4.dev3}/sqlmesh/core/linter/rule.py

@@ -39,6 +39,22 @@ class Range:
     end: Position
 
 
+@dataclass(frozen=True)
+class TextEdit:
+    """A text edit to apply to a file."""
+
+    range: Range
+    new_text: str
+
+
+@dataclass(frozen=True)
+class Fix:
+    """A fix that can be applied to resolve a rule violation."""
+
+    title: str
+    edits: t.List[TextEdit]
+
+
 class _Rule(abc.ABCMeta):
     def __new__(cls: Type[_Rule], clsname: str, bases: t.Tuple, attrs: t.Dict) -> _Rule:
         attrs["name"] = clsname.lower()
@@ -66,10 +82,14 @@ class Rule(abc.ABC, metaclass=_Rule):
         self,
         violation_msg: t.Optional[str] = None,
         violation_range: t.Optional[Range] = None,
+        fixes: t.Optional[t.List[Fix]] = None,
     ) -> RuleViolation:
         """Create a RuleViolation instance for this rule"""
         return RuleViolation(
-            rule=self, violation_msg=violation_msg or self.summary, violation_range=violation_range
+            rule=self,
+            violation_msg=violation_msg or self.summary,
+            violation_range=violation_range,
+            fixes=fixes,
         )
 
     def get_definition_location(self) -> RuleLocation:
@@ -103,11 +123,16 @@ class Rule(abc.ABC, metaclass=_Rule):
 
 class RuleViolation:
     def __init__(
-        self, rule: Rule, violation_msg: str, violation_range: t.Optional[Range] = None
+        self,
+        rule: Rule,
+        violation_msg: str,
+        violation_range: t.Optional[Range] = None,
+        fixes: t.Optional[t.List[Fix]] = None,
     ) -> None:
         self.rule = rule
         self.violation_msg = violation_msg
         self.violation_range = violation_range
+        self.fixes = fixes or []
 
     def __repr__(self) -> str:
         return f"{self.rule.name}: {self.violation_msg}"

{sqlmesh-0.196.4.dev2 → sqlmesh-0.196.4.dev3}/sqlmesh/core/linter/rules/builtin.py

@@ -8,7 +8,7 @@ from sqlglot.expressions import Star
 from sqlglot.helper import subclasses
 
 from sqlmesh.core.linter.helpers import TokenPositionDetails
-from sqlmesh.core.linter.rule import Rule, RuleViolation, Range
+from sqlmesh.core.linter.rule import Rule, RuleViolation, Range, Fix, TextEdit
 from sqlmesh.core.linter.definition import RuleSet
 from sqlmesh.core.model import Model, SqlModel
 
@@ -22,7 +22,8 @@ class NoSelectStar(Rule):
             return None
         if model.query.is_star:
             violation_range = self._get_range(model)
-            return self.violation(violation_range=violation_range)
+            fixes = self._create_fixes(model, violation_range)
+            return self.violation(violation_range=violation_range, fixes=fixes)
         return None
 
     def _get_range(self, model: SqlModel) -> t.Optional[Range]:
@@ -37,6 +38,28 @@ class NoSelectStar(Rule):
 
         return None
 
+    def _create_fixes(
+        self, model: SqlModel, violation_range: t.Optional[Range]
+    ) -> t.Optional[t.List[Fix]]:
+        """Create fixes for the SELECT * violation."""
+        if not violation_range:
+            return None
+        columns = model.columns_to_types
+        if not columns:
+            return None
+        new_text = ", ".join(columns.keys())
+        return [
+            Fix(
+                title="Replace SELECT * with explicit column list",
+                edits=[
+                    TextEdit(
+                        range=violation_range,
+                        new_text=new_text,
+                    )
+                ],
+            )
+        ]
+
 
 class InvalidSelectStarExpansion(Rule):
     def check_model(self, model: Model) -> t.Optional[RuleViolation]:

sqlmesh-0.196.4.dev3/sqlmesh/lsp/context.py

@@ -0,0 +1,346 @@
+from dataclasses import dataclass
+from pathlib import Path
+import uuid
+from sqlmesh.core.context import Context
+import typing as t
+
+from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.core.linter.definition import AnnotatedRuleViolation
+from sqlmesh.lsp.custom import ModelForRendering
+from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
+from sqlmesh.lsp.uri import URI
+from lsprotocol import types
+
+
+@dataclass
+class ModelTarget:
+    """Information about models in a file."""
+
+    names: t.List[str]
+
+
+@dataclass
+class AuditTarget:
+    """Information about standalone audits in a file."""
+
+    name: str
+
+
+class LSPContext:
+    """
+    A context that is used for linting. It contains the context and a reverse map of file uri to
+    model names and standalone audit names.
+    """
+
+    map: t.Dict[Path, t.Union[ModelTarget, AuditTarget]]
+    _render_cache: t.Dict[Path, t.List[RenderModelEntry]]
+    _lint_cache: t.Dict[Path, t.List[AnnotatedRuleViolation]]
+    _version_id: str
+    """
+    This is a version ID for the context. It is used to track changes to the context. It can be used to 
+    return a version number to the LSP client.
+    """
+
+    def __init__(self, context: Context) -> None:
+        self._version_id = str(uuid.uuid4())
+        self.context = context
+        self._render_cache = {}
+        self._lint_cache = {}
+
+        # Add models to the map
+        model_map: t.Dict[Path, ModelTarget] = {}
+        for model in context.models.values():
+            if model._path is not None:
+                uri = model._path
+                if uri in model_map:
+                    model_map[uri].names.append(model.name)
+                else:
+                    model_map[uri] = ModelTarget(names=[model.name])
+
+        # Add standalone audits to the map
+        audit_map: t.Dict[Path, AuditTarget] = {}
+        for audit in context.standalone_audits.values():
+            if audit._path is not None:
+                uri = audit._path
+                if uri not in audit_map:
+                    audit_map[uri] = AuditTarget(name=audit.name)
+
+        self.map: t.Dict[Path, t.Union[ModelTarget, AuditTarget]] = {
+            **model_map,
+            **audit_map,
+        }
+
+    @property
+    def version_id(self) -> str:
+        """Get the version ID for the context."""
+        return self._version_id
+
+    def render_model(self, uri: URI) -> t.List[RenderModelEntry]:
+        """Get rendered models for a file, using cache when available.
+
+        Args:
+            uri: The URI of the file to render.
+
+        Returns:
+            List of rendered model entries.
+        """
+        path = uri.to_path()
+
+        # Check cache first
+        if path in self._render_cache:
+            return self._render_cache[path]
+
+        # If not cached, render and cache
+        entries: t.List[RenderModelEntry] = []
+        target = self.map.get(path)
+
+        if isinstance(target, AuditTarget):
+            audit = self.context.standalone_audits[target.name]
+            definition = audit.render_definition(
+                include_python=False,
+                render_query=True,
+            )
+            rendered_query = [
+                render.sql(dialect=audit.dialect, pretty=True) for render in definition
+            ]
+            entry = RenderModelEntry(
+                name=audit.name,
+                fqn=audit.fqn,
+                description=audit.description,
+                rendered_query="\n\n".join(rendered_query),
+            )
+            entries.append(entry)
+
+        elif isinstance(target, ModelTarget):
+            for name in target.names:
+                model = self.context.get_model(name)
+                if isinstance(model, SqlModel):
+                    rendered_query = [
+                        render.sql(dialect=model.dialect, pretty=True)
+                        for render in model.render_definition(
+                            include_python=False,
+                            render_query=True,
+                        )
+                    ]
+                    entry = RenderModelEntry(
+                        name=model.name,
+                        fqn=model.fqn,
+                        description=model.description,
+                        rendered_query="\n\n".join(rendered_query),
+                    )
+                    entries.append(entry)
+
+        # Store in cache
+        self._render_cache[path] = entries
+        return entries
+
+    def lint_model(self, uri: URI) -> t.List[AnnotatedRuleViolation]:
+        """Get lint diagnostics for a model, using cache when available.
+
+        Args:
+            uri: The URI of the file to lint.
+
+        Returns:
+            List of annotated rule violations.
+        """
+        path = uri.to_path()
+
+        # Check cache first
+        if path in self._lint_cache:
+            return self._lint_cache[path]
+
+        # If not cached, lint and cache
+        target = self.map.get(path)
+        if target is None or not isinstance(target, ModelTarget):
+            return []
+
+        diagnostics = self.context.lint_models(
+            target.names,
+            raise_on_error=False,
+        )
+
+        # Store in cache
+        self._lint_cache[path] = diagnostics
+        return diagnostics
+
+    def get_code_actions(
+        self, uri: URI, params: types.CodeActionParams
+    ) -> t.Optional[t.List[t.Union[types.Command, types.CodeAction]]]:
+        """Get code actions for a file."""
+
+        # Get the violations (which contain the fixes)
+        violations = self.lint_model(uri)
+
+        # Convert violations to a map for quick lookup
+        # Use a hashable representation of Range as the key
+        violation_map: t.Dict[
+            t.Tuple[str, t.Tuple[int, int, int, int]], AnnotatedRuleViolation
+        ] = {}
+        for violation in violations:
+            if violation.violation_range:
+                lsp_diagnostic = self.diagnostic_to_lsp_diagnostic(violation)
+                if lsp_diagnostic:
+                    # Create a hashable key from the diagnostic message and range
+                    key = (
+                        lsp_diagnostic.message,
+                        (
+                            lsp_diagnostic.range.start.line,
+                            lsp_diagnostic.range.start.character,
+                            lsp_diagnostic.range.end.line,
+                            lsp_diagnostic.range.end.character,
+                        ),
+                    )
+                    violation_map[key] = violation
+
+        # Get diagnostics in the requested range
+        diagnostics = params.context.diagnostics if params.context else []
+
+        code_actions: t.List[t.Union[types.Command, types.CodeAction]] = []
+
+        for diagnostic in diagnostics:
+            # Find the corresponding violation
+            key = (
+                diagnostic.message,
+                (
+                    diagnostic.range.start.line,
+                    diagnostic.range.start.character,
+                    diagnostic.range.end.line,
+                    diagnostic.range.end.character,
+                ),
+            )
+            found_violation = violation_map.get(key)
+
+            if found_violation is not None and found_violation.fixes:
+                # Create code actions for each fix
+                for fix in found_violation.fixes:
+                    # Convert our Fix to LSP TextEdits
+                    text_edits = []
+                    for edit in fix.edits:
+                        text_edits.append(
+                            types.TextEdit(
+                                range=types.Range(
+                                    start=types.Position(
+                                        line=edit.range.start.line,
+                                        character=edit.range.start.character,
+                                    ),
+                                    end=types.Position(
+                                        line=edit.range.end.line,
+                                        character=edit.range.end.character,
+                                    ),
+                                ),
+                                new_text=edit.new_text,
+                            )
+                        )
+
+                    # Create the code action
+                    code_action = types.CodeAction(
+                        title=fix.title,
+                        kind=types.CodeActionKind.QuickFix,
+                        diagnostics=[diagnostic],
+                        edit=types.WorkspaceEdit(changes={params.text_document.uri: text_edits}),
+                    )
+                    code_actions.append(code_action)
+
+        return code_actions if code_actions else None
+
+    def list_of_models_for_rendering(self) -> t.List[ModelForRendering]:
+        """Get a list of models for rendering.
+
+        Returns:
+            List of ModelForRendering objects.
+        """
+        return [
+            ModelForRendering(
+                name=model.name,
+                fqn=model.fqn,
+                description=model.description,
+                uri=URI.from_path(model._path).value,
+            )
+            for model in self.context.models.values()
+            if isinstance(model, SqlModel) and model._path is not None
+        ] + [
+            ModelForRendering(
+                name=audit.name,
+                fqn=audit.fqn,
+                description=audit.description,
+                uri=URI.from_path(audit._path).value,
+            )
+            for audit in self.context.standalone_audits.values()
+            if audit._path is not None
+        ]
+
+    @staticmethod
+    def get_completions(
+        self: t.Optional["LSPContext"] = None,
+        uri: t.Optional[URI] = None,
+        file_content: t.Optional[str] = None,
+    ) -> AllModelsResponse:
+        """Get completion suggestions for a file"""
+        from sqlmesh.lsp.completions import get_sql_completions
+
+        return get_sql_completions(self, uri, file_content)
+
+    @staticmethod
+    def diagnostics_to_lsp_diagnostics(
+        diagnostics: t.List[AnnotatedRuleViolation],
+    ) -> t.List[types.Diagnostic]:
+        """
+        Converts a list of AnnotatedRuleViolations to a list of LSP diagnostics. It will remove duplicates based on the message and range.
+        """
+        lsp_diagnostics = {}
+        for diagnostic in diagnostics:
+            lsp_diagnostic = LSPContext.diagnostic_to_lsp_diagnostic(diagnostic)
+            if lsp_diagnostic is not None:
+                # Create a unique key combining message and range
+                diagnostic_key = (
+                    lsp_diagnostic.message,
+                    lsp_diagnostic.range.start.line,
+                    lsp_diagnostic.range.start.character,
+                    lsp_diagnostic.range.end.line,
+                    lsp_diagnostic.range.end.character,
+                )
+                if diagnostic_key not in lsp_diagnostics:
+                    lsp_diagnostics[diagnostic_key] = lsp_diagnostic
+        return list(lsp_diagnostics.values())
+
+    @staticmethod
+    def diagnostic_to_lsp_diagnostic(
+        diagnostic: AnnotatedRuleViolation,
+    ) -> t.Optional[types.Diagnostic]:
+        if diagnostic.model._path is None:
+            return None
+        if not diagnostic.violation_range:
+            with open(diagnostic.model._path, "r", encoding="utf-8") as file:
+                lines = file.readlines()
+            diagnostic_range = types.Range(
+                start=types.Position(line=0, character=0),
+                end=types.Position(line=len(lines) - 1, character=len(lines[-1])),
+            )
+        else:
+            diagnostic_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()
+        rule_uri_wihout_extension = URI.from_path(rule_location.file_path)
+        rule_uri = f"{rule_uri_wihout_extension.value}#L{rule_location.start_line}"
+
+        # Use URI format to create a link for "related information"
+        return types.Diagnostic(
+            range=diagnostic_range,
+            message=diagnostic.violation_msg,
+            severity=types.DiagnosticSeverity.Error
+            if diagnostic.violation_type == "error"
+            else types.DiagnosticSeverity.Warning,
+            source="sqlmesh",
+            code=diagnostic.rule.name,
+            code_description=types.CodeDescription(href=rule_uri),
+        )