sqlmesh 0.184.2.dev2__py3-none-any.whl → 0.184.3__py3-none-any.whl

sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.184.2.dev2'
-__version_tuple__ = version_tuple = (0, 184, 2, 'dev2')
+__version__ = version = '0.184.3'
+__version_tuple__ = version_tuple = (0, 184, 3)

sqlmesh/cli/main.py

@@ -906,6 +906,11 @@ def create_external_models(obj: Context, **kwargs: t.Any) -> None:
     is_flag=True,
     help="Disable the check for a primary key (grain) that is missing or is not unique.",
 )
+@click.option(
+    "--warn-grain-check",
+    is_flag=True,
+    help="Warn if any selected model is missing a grain, and compute diffs for the remaining models.",
+)
 @click.option(
     "--temp-schema",
     type=str,

sqlmesh/core/config/connection.py

@@ -4,6 +4,7 @@ import abc
 import base64
 import logging
 import os
+import importlib
 import pathlib
 import re
 import typing as t
@@ -24,6 +25,7 @@ from sqlmesh.core.config.common import (
 )
 from sqlmesh.core.engine_adapter.shared import CatalogSupport
 from sqlmesh.core.engine_adapter import EngineAdapter
+from sqlmesh.utils import str_to_bool
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import (
     ValidationInfo,
@@ -49,6 +51,30 @@ FORBIDDEN_STATE_SYNC_ENGINES = {
 MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)")
 
 
+def _get_engine_import_validator(
+    import_name: str, engine_type: str, extra_name: t.Optional[str] = None
+) -> t.Callable:
+    extra_name = extra_name or engine_type
+
+    @model_validator(mode="before")
+    def validate(cls: t.Any, data: t.Any) -> t.Any:
+        check_import = (
+            str_to_bool(str(data.pop("check_import", True))) if isinstance(data, dict) else True
+        )
+        if not check_import:
+            return data
+        try:
+            importlib.import_module(import_name)
+        except ImportError:
+            raise ConfigError(
+                f"Failed to import the '{engine_type}' engine library. Please run `pip install \"sqlmesh[{extra_name}]\"`."
+            )
+
+        return data
+
+    return validate
+
+
 class ConnectionConfig(abc.ABC, BaseConfig):
     type_: str
     concurrent_tasks: int
@@ -428,6 +454,7 @@ class SnowflakeConnectionConfig(ConnectionConfig):
     type_: t.Literal["snowflake"] = Field(alias="type", default="snowflake")
 
     _concurrent_tasks_validator = concurrent_tasks_validator
+    _engine_import_validator = _get_engine_import_validator("snowflake", "snowflake")
 
     @model_validator(mode="before")
     def _validate_authenticator(cls, data: t.Any) -> t.Any:
@@ -621,6 +648,7 @@ class DatabricksConnectionConfig(ConnectionConfig):
 
     _concurrent_tasks_validator = concurrent_tasks_validator
     _http_headers_validator = http_headers_validator
+    _engine_import_validator = _get_engine_import_validator("databricks", "databricks")
 
     @model_validator(mode="before")
     def _databricks_connect_validator(cls, data: t.Any) -> t.Any:
@@ -873,6 +901,8 @@ class BigQueryConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["bigquery"] = Field(alias="type", default="bigquery")
 
+    _engine_import_validator = _get_engine_import_validator("google.cloud.bigquery", "bigquery")
+
     @field_validator("execution_project")
     def validate_execution_project(
         cls,
@@ -1015,6 +1045,10 @@ class GCPPostgresConnectionConfig(ConnectionConfig):
     register_comments: bool = True
     pre_ping: bool = True
 
+    _engine_import_validator = _get_engine_import_validator(
+        "google.cloud.sql", "gcp_postgres", "gcppostgres"
+    )
+
     @model_validator(mode="before")
     def _validate_auth_method(cls, data: t.Any) -> t.Any:
         if not isinstance(data, dict):
@@ -1142,6 +1176,8 @@ class RedshiftConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["redshift"] = Field(alias="type", default="redshift")
 
+    _engine_import_validator = _get_engine_import_validator("redshift_connector", "redshift")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1201,6 +1237,8 @@ class PostgresConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["postgres"] = Field(alias="type", default="postgres")
 
+    _engine_import_validator = _get_engine_import_validator("psycopg2", "postgres")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1252,6 +1290,8 @@ class MySQLConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["mysql"] = Field(alias="type", default="mysql")
 
+    _engine_import_validator = _get_engine_import_validator("pymysql", "mysql")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         connection_keys = {
@@ -1302,6 +1342,8 @@ class MSSQLConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["mssql"] = Field(alias="type", default="mssql")
 
+    _engine_import_validator = _get_engine_import_validator("pymssql", "mssql")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1357,6 +1399,8 @@ class SparkConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["spark"] = Field(alias="type", default="spark")
 
+    _engine_import_validator = _get_engine_import_validator("pyspark", "spark")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1473,6 +1517,8 @@ class TrinoConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["trino"] = Field(alias="type", default="trino")
 
+    _engine_import_validator = _get_engine_import_validator("trino", "trino")
+
     @field_validator("schema_location_mapping", mode="before")
     @classmethod
     def _validate_regex_keys(
@@ -1623,6 +1669,8 @@ class ClickhouseConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["clickhouse"] = Field(alias="type", default="clickhouse")
 
+    _engine_import_validator = _get_engine_import_validator("clickhouse_connect", "clickhouse")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         kwargs = {
@@ -1727,6 +1775,8 @@ class AthenaConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["athena"] = Field(alias="type", default="athena")
 
+    _engine_import_validator = _get_engine_import_validator("pyathena", "athena")
+
     @model_validator(mode="after")
     def _root_validator(self) -> Self:
         work_group = self.work_group
@@ -1793,6 +1843,8 @@ class RisingwaveConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["risingwave"] = Field(alias="type", default="risingwave")
 
+    _engine_import_validator = _get_engine_import_validator("psycopg2", "risingwave")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {

sqlmesh/core/console.py

@@ -1161,7 +1161,7 @@ class TerminalConsole(Console):
         default_catalog: t.Optional[str],
     ) -> None:
         """Indicates that a new snapshot promotion progress has begun."""
-        if self.promotion_progress is None:
+        if snapshots and self.promotion_progress is None:
             self.promotion_progress = make_progress_bar(
                 "Updating virtual layer ", self.console, justify="left"
             )
@@ -1782,10 +1782,13 @@ class TerminalConsole(Console):
         """Get the user's change category for the directly modified models."""
         plan = plan_builder.build()
 
-        self.show_environment_difference_summary(
-            plan.context_diff,
-            no_diff=no_diff,
-        )
+        if plan.restatements:
+            self._print("\n[bold]Restating models\n")
+        else:
+            self.show_environment_difference_summary(
+                plan.context_diff,
+                no_diff=no_diff,
+            )
 
         if plan.context_diff.has_changes:
             self.show_model_difference_summary(
@@ -3458,7 +3461,8 @@ class DebuggerTerminalConsole(TerminalConsole):
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
-        self._write(f"Starting promotion for {len(snapshots)} snapshots")
+        if snapshots:
+            self._write(f"Starting promotion for {len(snapshots)} snapshots")
 
     def update_promotion_progress(self, snapshot: SnapshotInfoLike, promoted: bool) -> None:
         self._write(f"Promoting {snapshot.name}")

sqlmesh/core/context.py

@@ -1629,6 +1629,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         show_sample: bool = True,
         decimals: int = 3,
         skip_grain_check: bool = False,
+        warn_grain_check: bool = False,
         temp_schema: t.Optional[str] = None,
     ) -> t.List[TableDiff]:
         """Show a diff between two tables.
@@ -1703,22 +1704,29 @@ class GenericContext(BaseContext, t.Generic[C]):
                             target_env.naming_info, adapter.dialect
                         )
                         model_on = on or model.on
-                        models_to_diff.append((model, adapter, source, target, model_on))
                         if not model_on:
                             models_without_grain.append(model)
+                        else:
+                            models_to_diff.append((model, adapter, source, target, model_on))
+
+            if models_without_grain:
+                model_names = "\n".join(
+                    f"─ {model.name} \n  at '{model._path}'" for model in models_without_grain
+                )
+                message = (
+                    "SQLMesh doesn't know how to join the tables for the following models:\n"
+                    f"{model_names}\n\n"
+                    "Please specify a `grain` in each model definition. It must be unique and not null."
+                )
+                if warn_grain_check:
+                    self.console.log_warning(message)
+                else:
+                    raise SQLMeshError(message)
 
             if models_to_diff:
                 self.console.show_table_diff_details(
                     [model[0].name for model in models_to_diff],
                 )
-                if models_without_grain:
-                    model_names = "\n".join(
-                        f"─ {model.name} \n  at '{model._path}'" for model in models_without_grain
-                    )
-                    raise SQLMeshError(
-                        f"SQLMesh doesn't know how to join the tables for the following models:\n{model_names}\n"
-                        "\nPlease specify the `grain` in each model definition. Must be unique and not null."
-                    )
 
                 self.console.start_table_diff_progress(len(models_to_diff))
                 try:

sqlmesh/core/engine_adapter/athena.py

@@ -8,7 +8,7 @@ from sqlmesh.utils.aws import validate_s3_uri, parse_s3_uri
 from sqlmesh.core.engine_adapter.mixins import PandasNativeFetchDFSupportMixin, RowDiffMixin
 from sqlmesh.core.engine_adapter.trino import TrinoEngineAdapter
 from sqlmesh.core.node import IntervalUnit
-import os
+import posixpath
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.core.engine_adapter.shared import (
     CatalogSupport,
@@ -403,11 +403,13 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin, RowDiffMixin):
 
         elif self.s3_warehouse_location:
             # If the user has set `s3_warehouse_location` in the connection config, the base URI is <s3_warehouse_location>/<catalog>/<schema>/
-            base_uri = os.path.join(self.s3_warehouse_location, table.catalog or "", table.db or "")
+            base_uri = posixpath.join(
+                self.s3_warehouse_location, table.catalog or "", table.db or ""
+            )
         else:
             return None
 
-        full_uri = validate_s3_uri(os.path.join(base_uri, table.text("this") or ""), base=True)
+        full_uri = validate_s3_uri(posixpath.join(base_uri, table.text("this") or ""), base=True)
         return exp.LocationProperty(this=exp.Literal.string(full_uri))
 
     def _find_matching_columns(

sqlmesh/core/engine_adapter/snowflake.py

@@ -68,6 +68,7 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
         },
     )
     MANAGED_TABLE_KIND = "DYNAMIC TABLE"
+    SNOWPARK = "snowpark"
 
     @contextlib.contextmanager
     def session(self, properties: SessionProperties) -> t.Iterator[None]:
@@ -104,9 +105,16 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
     @property
     def snowpark(self) -> t.Optional[SnowparkSession]:
         if snowpark:
-            return snowpark.Session.builder.configs(
-                {"connection": self._connection_pool.get()}
-            ).getOrCreate()
+            if not self._connection_pool.get_attribute(self.SNOWPARK):
+                # Snowpark sessions are not thread safe so we create a session per thread to prevent them from interfering with each other
+                # The sessions are cleaned up when close() is called
+                new_session = snowpark.Session.builder.configs(
+                    {"connection": self._connection_pool.get()}
+                ).create()
+                self._connection_pool.set_attribute(self.SNOWPARK, new_session)
+
+            return self._connection_pool.get_attribute(self.SNOWPARK)
+
         return None
 
     @property
@@ -584,3 +592,10 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
             return columns_to_types_from_dtypes(query_or_df.sample(n=1).to_pandas().dtypes.items())
 
         return super()._columns_to_types(query_or_df, columns_to_types)
+
+    def close(self) -> t.Any:
+        if snowpark_session := self._connection_pool.get_attribute(self.SNOWPARK):
+            snowpark_session.close()  # type: ignore
+            self._connection_pool.set_attribute(self.SNOWPARK, None)
+
+        return super().close()

sqlmesh/core/macros.py

@@ -1090,10 +1090,10 @@ def pivot(
     column: SQL,
     values: t.List[SQL],
     alias: bool = True,
-    agg: SQL = SQL("SUM"),
-    cmp: SQL = SQL("="),
-    prefix: SQL = SQL(""),
-    suffix: SQL = SQL(""),
+    agg: exp.Expression = exp.Literal.string("SUM"),
+    cmp: exp.Expression = exp.Literal.string("="),
+    prefix: exp.Expression = exp.Literal.string(""),
+    suffix: exp.Expression = exp.Literal.string(""),
     then_value: SQL = SQL("1"),
     else_value: SQL = SQL("0"),
     quote: bool = True,
@@ -1107,18 +1107,26 @@ def pivot(
         >>> sql = "SELECT date_day, @PIVOT(status, ['cancelled', 'completed']) FROM rides GROUP BY 1"
         >>> MacroEvaluator().transform(parse_one(sql)).sql()
         'SELECT date_day, SUM(CASE WHEN status = \\'cancelled\\' THEN 1 ELSE 0 END) AS "\\'cancelled\\'", SUM(CASE WHEN status = \\'completed\\' THEN 1 ELSE 0 END) AS "\\'completed\\'" FROM rides GROUP BY 1'
+        >>> sql = "SELECT @PIVOT(a, ['v'], then_value := tv, suffix := '_sfx', quote := FALSE)"
+        >>> MacroEvaluator(dialect="bigquery").transform(parse_one(sql)).sql("bigquery")
+        "SELECT SUM(CASE WHEN a = 'v' THEN tv ELSE 0 END) AS `v_sfx`"
     """
     aggregates: t.List[exp.Expression] = []
     for value in values:
-        proj = f"{agg}("
+        proj = f"{agg.name}("
         if distinct:
             proj += "DISTINCT "
 
-        proj += f"CASE WHEN {column} {cmp} {value} THEN {then_value} ELSE {else_value} END) "
+        proj += f"CASE WHEN {column} {cmp.name} {value} THEN {then_value} ELSE {else_value} END) "
         node = evaluator.parse_one(proj)
 
         if alias:
-            node = node.as_(f"{prefix}{value}{suffix}", quoted=quote, copy=False)
+            node = node.as_(
+                f"{prefix.name}{value}{suffix.name}",
+                quoted=quote,
+                copy=False,
+                dialect=evaluator.dialect,
+            )
 
         aggregates.append(node)
 

sqlmesh/core/model/cache.py

@@ -15,6 +15,7 @@ from sqlmesh.core import constants as c
 from sqlmesh.core.model.definition import ExternalModel, Model, SqlModel, _Model
 from sqlmesh.utils.cache import FileCache
 from sqlmesh.utils.hashing import crc32
+from sqlmesh.utils.windows import IS_WINDOWS
 
 from dataclasses import dataclass
 
@@ -149,8 +150,10 @@ class OptimizedQueryCache:
 
 
 def optimized_query_cache_pool(optimized_query_cache: OptimizedQueryCache) -> ProcessPoolExecutor:
+    # fork doesnt work on Windows. ref: https://docs.python.org/3/library/multiprocessing.html#multiprocessing-start-methods
+    context_type = "spawn" if IS_WINDOWS else "fork"
     return ProcessPoolExecutor(
-        mp_context=mp.get_context("fork"),
+        mp_context=mp.get_context(context_type),
         initializer=_init_optimized_query_cache,
         initargs=(optimized_query_cache,),
         max_workers=c.MAX_FORK_WORKERS,

sqlmesh/core/plan/evaluator.py

@@ -133,7 +133,9 @@ class BuiltInPlanEvaluator(PlanEvaluator):
             push_completion_status = self._push(plan, snapshots, deployability_index_for_creation)
             if push_completion_status.is_nothing_to_do:
                 self.console.log_status_update(
-                    "\n[green]SKIP: No physical layer updates to perform[/green]\n"
+                    ""
+                    if plan.restatements
+                    else "\n[green]SKIP: No physical layer updates to perform[/green]\n"
                 )
             update_intervals_for_new_snapshots(plan.new_snapshots, self.state_sync)
             self._restate(plan, snapshots_by_name)

sqlmesh/integrations/github/cicd/controller.py

@@ -457,10 +457,13 @@ class GithubController:
         try:
             # Clear out any output that might exist from prior steps
             self._console.clear_captured_outputs()
-            self._console.show_environment_difference_summary(
-                context_diff=plan.context_diff,
-                no_diff=False,
-            )
+            if plan.restatements:
+                self._console._print("\n**Restating models**\n")
+            else:
+                self._console.show_environment_difference_summary(
+                    context_diff=plan.context_diff,
+                    no_diff=False,
+                )
             if plan.context_diff.has_changes:
                 self._console.show_model_difference_summary(
                     context_diff=plan.context_diff,

sqlmesh/lsp/main.py

@@ -92,9 +92,7 @@ class SQLMeshLanguageServer:
                 return get_sql_completions(None, uri)
 
         @self.server.feature(API_FEATURE)
-        def api(
-            ls: LanguageServer, request: ApiRequest
-        ) -> t.Union[ApiResponseGetModels, ApiResponseGetLineage, ApiResponseGetColumnLineage]:
+        def api(ls: LanguageServer, request: ApiRequest) -> t.Dict[str, t.Any]:
             ls.log_trace(f"API request: {request}")
             if self.lsp_context is None:
                 raise RuntimeError("No context found")
@@ -105,7 +103,9 @@ class SQLMeshLanguageServer:
             if request.method == "GET":
                 if path_parts == ["api", "models"]:
                     # /api/models
-                    return ApiResponseGetModels(data=get_models(self.lsp_context.context))
+                    return ApiResponseGetModels(
+                        data=get_models(self.lsp_context.context)
+                    ).model_dump(mode="json")
 
                 if path_parts[:2] == ["api", "lineage"]:
                     if len(path_parts) == 3:
@@ -113,7 +113,7 @@ class SQLMeshLanguageServer:
                         model_name = urllib.parse.unquote(path_parts[2])
                         lineage = model_lineage(model_name, self.lsp_context.context)
                         non_set_lineage = {k: v for k, v in lineage.items() if v is not None}
-                        return ApiResponseGetLineage(data=non_set_lineage)
+                        return ApiResponseGetLineage(data=non_set_lineage).model_dump(mode="json")
 
                     if len(path_parts) == 4:
                         # /api/lineage/{model}/{column}
@@ -123,7 +123,9 @@ class SQLMeshLanguageServer:
                         column_lineage_response = column_lineage(
                             model_name, column, False, self.lsp_context.context
                         )
-                        return ApiResponseGetColumnLineage(data=column_lineage_response)
+                        return ApiResponseGetColumnLineage(data=column_lineage_response).model_dump(
+                            mode="json"
+                        )
 
             raise NotImplementedError(f"API request not implemented: {request.url}")
 

sqlmesh/magics.py

@@ -702,6 +702,11 @@ class SQLMeshMagics(Magics):
         action="store_true",
         help="Disable the check for a primary key (grain) that is missing or is not unique.",
     )
+    @argument(
+        "--warn-grain-check",
+        action="store_true",
+        help="Warn if any selected model is missing a grain, and compute diffs for the remaining models.",
+    )
     @line_magic
     @pass_sqlmesh_context
     def table_diff(self, context: Context, line: str) -> None:
@@ -723,6 +728,7 @@ class SQLMeshMagics(Magics):
             show_sample=args.show_sample,
             decimals=args.decimals,
             skip_grain_check=args.skip_grain_check,
+            warn_grain_check=args.warn_grain_check,
         )
 
     @magic_arguments()

sqlmesh/utils/cache.py

@@ -12,6 +12,7 @@ from sqlglot import __version__ as SQLGLOT_VERSION
 from sqlmesh.utils import sanitize_name
 from sqlmesh.utils.date import to_datetime
 from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils.windows import IS_WINDOWS, fix_windows_path
 
 logger = logging.getLogger(__name__)
 
@@ -132,4 +133,8 @@ class FileCache(t.Generic[T]):
 
     def _cache_entry_path(self, name: str, entry_id: str = "") -> Path:
         entry_file_name = "__".join(p for p in (self._cache_version, name, entry_id) if p)
-        return self._path / sanitize_name(entry_file_name)
+        full_path = self._path / sanitize_name(entry_file_name)
+        if IS_WINDOWS:
+            # handle paths longer than 260 chars
+            full_path = fix_windows_path(full_path)
+        return full_path

sqlmesh/utils/windows.py

@@ -0,0 +1,14 @@
+import platform
+from pathlib import Path
+
+IS_WINDOWS = platform.system() == "Windows"
+
+
+def fix_windows_path(path: Path) -> Path:
+    """
+    Windows paths are limited to 260 characters: https://learn.microsoft.com/en-us/windows/win32/fileio/maximum-file-path-limitation
+    Users can change this by updating a registry entry but we cant rely on that.
+    We can quite commonly generate a cache file path that exceeds 260 characters which causes a FileNotFound error.
+    If we prefix the path with "\\?\" then we can have paths up to 32,767 characters
+    """
+    return Path("\\\\?\\" + str(path.absolute()))

{sqlmesh-0.184.2.dev2.dist-info → sqlmesh-0.184.3.dist-info}/METADATA

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