sqlmesh 0.196.3.dev4__py3-none-any.whl → 0.196.4__py3-none-any.whl

sqlmesh/_version.py

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

sqlmesh/cli/main.py

@@ -4,25 +4,26 @@ import logging
 import os
 import sys
 import typing as t
+from pathlib import Path
 
 import click
+
 from sqlmesh import configure_logging, remove_excess_logs
 from sqlmesh.cli import error_handler
 from sqlmesh.cli import options as opt
 from sqlmesh.cli.project_init import (
+    InitCliMode,
     ProjectTemplate,
     init_example_project,
-    InitCliMode,
     interactive_init,
 )
 from sqlmesh.core.analytics import cli_analytics
-from sqlmesh.core.console import configure_console, get_console
-from sqlmesh.utils import Verbosity
 from sqlmesh.core.config import load_configs
+from sqlmesh.core.console import configure_console, get_console
 from sqlmesh.core.context import Context
+from sqlmesh.utils import Verbosity
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import MissingDependencyError, SQLMeshError
-from pathlib import Path
 
 logger = logging.getLogger(__name__)
 
@@ -355,6 +356,7 @@ def evaluate(
     "--append-newline",
     is_flag=True,
     help="Include a newline at the end of each file.",
+    default=None,
 )
 @opt.format_options
 @click.pass_context
@@ -805,7 +807,11 @@ def check_intervals(
     context = ctx.obj
     context.console.show_intervals(
         context.check_intervals(
-            environment, no_signals=no_signals, select_models=select_model, start=start, end=end
+            environment,
+            no_signals=no_signals,
+            select_models=select_model,
+            start=start,
+            end=end,
         )
     )
 
@@ -1107,7 +1113,10 @@ def clean(obj: Context) -> None:
 @error_handler
 @cli_analytics
 def table_name(
-    obj: Context, model_name: str, environment: t.Optional[str] = None, prod: bool = False
+    obj: Context,
+    model_name: str,
+    environment: t.Optional[str] = None,
+    prod: bool = False,
 ) -> None:
     """Prints the name of the physical table for the given model."""
     print(obj.table_name(model_name, environment, prod))

sqlmesh/core/engine_adapter/mssql.py

@@ -6,11 +6,13 @@ import typing as t
 
 from sqlglot import exp
 
-from sqlmesh.core.dialect import to_schema
+from sqlmesh.core.dialect import to_schema, add_table
 from sqlmesh.core.engine_adapter.base import (
     EngineAdapterWithIndexSupport,
     EngineAdapter,
     InsertOverwriteStrategy,
+    MERGE_SOURCE_ALIAS,
+    MERGE_TARGET_ALIAS,
 )
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
@@ -32,7 +34,7 @@ from sqlmesh.core.schema_diff import SchemaDiffer
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import SchemaName, TableName
-    from sqlmesh.core.engine_adapter._typing import DF, Query
+    from sqlmesh.core.engine_adapter._typing import DF, Query, QueryOrDF
 
 
 @set_catalog()
@@ -188,6 +190,87 @@ class MSSQLEngineAdapter(
                     )
         super().drop_schema(schema_name, ignore_if_not_exists=ignore_if_not_exists, cascade=False)
 
+    def merge(
+        self,
+        target_table: TableName,
+        source_table: QueryOrDF,
+        columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
+        unique_key: t.Sequence[exp.Expression],
+        when_matched: t.Optional[exp.Whens] = None,
+        merge_filter: t.Optional[exp.Expression] = None,
+    ) -> None:
+        source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
+            source_table, columns_to_types, target_table=target_table
+        )
+        columns_to_types = columns_to_types or self.columns(target_table)
+        on = exp.and_(
+            *(
+                add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS))
+                for part in unique_key
+            )
+        )
+        if merge_filter:
+            on = exp.and_(merge_filter, on)
+
+        if not when_matched:
+            match_condition = None
+            unique_key_names = [y.name for y in unique_key]
+            columns_to_types_no_keys = [c for c in columns_to_types if c not in unique_key_names]
+
+            target_columns_no_keys = [
+                exp.column(c, MERGE_TARGET_ALIAS) for c in columns_to_types_no_keys
+            ]
+            source_columns_no_keys = [
+                exp.column(c, MERGE_SOURCE_ALIAS) for c in columns_to_types_no_keys
+            ]
+
+            match_condition = exp.Exists(
+                this=exp.select(*target_columns_no_keys).except_(
+                    exp.select(*source_columns_no_keys)
+                )
+            )
+
+            match_expressions = [
+                exp.When(
+                    matched=True,
+                    source=False,
+                    condition=match_condition,
+                    then=exp.Update(
+                        expressions=[
+                            exp.column(col, MERGE_TARGET_ALIAS).eq(
+                                exp.column(col, MERGE_SOURCE_ALIAS)
+                            )
+                            for col in columns_to_types_no_keys
+                        ],
+                    ),
+                )
+            ]
+        else:
+            match_expressions = when_matched.copy().expressions
+
+        match_expressions.append(
+            exp.When(
+                matched=False,
+                source=False,
+                then=exp.Insert(
+                    this=exp.Tuple(expressions=[exp.column(col) for col in columns_to_types]),
+                    expression=exp.Tuple(
+                        expressions=[
+                            exp.column(col, MERGE_SOURCE_ALIAS) for col in columns_to_types
+                        ]
+                    ),
+                ),
+            )
+        )
+        for source_query in source_queries:
+            with source_query as query:
+                self._merge(
+                    target_table=target_table,
+                    query=query,
+                    on=on,
+                    whens=exp.Whens(expressions=match_expressions),
+                )
+
     def _convert_df_datetime(self, df: DF, columns_to_types: t.Dict[str, exp.DataType]) -> None:
         import pandas as pd
         from pandas.api.types import is_datetime64_any_dtype  # type: ignore

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/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/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/core/model/definition.py

@@ -2886,6 +2886,7 @@ META_FIELD_CONVERTER: t.Dict[str, t.Callable] = {
             for name, args in values
         )
     ),
+    "formatting": str,
 }
 
 

sqlmesh/core/state_sync/db/environment.py

@@ -172,13 +172,9 @@ class EnvironmentState:
         Returns:
             The list of environment summaries to remove.
         """
-
-        environment_summaries = self.get_environments_summary()
-        return [
-            env_summary
-            for env_summary in environment_summaries
-            if env_summary.expiration_ts is not None and env_summary.expiration_ts <= current_ts
-        ]
+        return self._fetch_environment_summaries(
+            where=self._create_expiration_filter_expr(current_ts)
+        )
 
     def delete_expired_environments(
         self, current_ts: t.Optional[int] = None
@@ -225,13 +221,7 @@ class EnvironmentState:
         Returns:
             A list of all environment summaries.
         """
-        return [
-            self._environment_summmary_from_row(row)
-            for row in fetchall(
-                self.engine_adapter,
-                self._environments_query(required_fields=list(EnvironmentSummary.all_fields())),
-            )
-        ]
+        return self._fetch_environment_summaries()
 
     def get_environment(
         self, environment: str, lock_for_update: bool = False
@@ -327,6 +317,20 @@ class EnvironmentState:
             expression=exp.Literal.number(current_ts),
         )
 
+    def _fetch_environment_summaries(
+        self, where: t.Optional[str | exp.Expression] = None
+    ) -> t.List[EnvironmentSummary]:
+        return [
+            self._environment_summmary_from_row(row)
+            for row in fetchall(
+                self.engine_adapter,
+                self._environments_query(
+                    where=where,
+                    required_fields=list(EnvironmentSummary.all_fields()),
+                ),
+            )
+        ]
+
 
 def _environment_to_df(environment: Environment) -> pd.DataFrame:
     import pandas as pd

sqlmesh/core/state_sync/db/facade.py

@@ -69,10 +69,6 @@ logger = logging.getLogger(__name__)
 T = t.TypeVar("T")
 
 
-if t.TYPE_CHECKING:
-    pass
-
-
 class EngineAdapterStateSync(StateSync):
     """Manages state of nodes and snapshot with an existing engine adapter.
 

sqlmesh/lsp/context.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 from pathlib import Path
+import uuid
 from sqlmesh.core.context import Context
 import typing as t
 
@@ -8,6 +9,7 @@ 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
@@ -33,8 +35,14 @@ class LSPContext:
     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 = {}
@@ -62,6 +70,11 @@ class LSPContext:
             **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.
 
@@ -150,6 +163,86 @@ class LSPContext:
         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.
 
@@ -186,3 +279,68 @@ class LSPContext:
         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),
+        )