Merge branch 'main' into vchan/fix-macro-dispatch

Author: themisvaltinos

docs/integrations/engines/trino.md

@@ -81,19 +81,21 @@ hive.metastore.glue.default-warehouse-dir=s3://my-bucket/
 
 ### Connection options
 
-| Option               | Description                                                                                                                                                               |  Type  | Required |
-|----------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------:|:--------:|
-| `type`               | Engine type name - must be `trino`                                                                                                                                        | string |    Y     |
-| `user`               | The username (of the account) to log in to your cluster. When connecting to Starburst Galaxy clusters, you must include the role of the user as a suffix to the username. | string |    Y     |
-| `host`               | The hostname of your cluster. Don't include the `http://` or `https://` prefix.                                                                                           | string |    Y     |
-| `catalog`            | The name of a catalog in your cluster.                                                                                                                                    | string |    Y     |
-| `http_scheme`        | The HTTP scheme to use when connecting to your cluster. By default, it's `https` and can only be `http` for no-auth or basic auth.                                        | string |    N     |
-| `port`               | The port to connect to your cluster. By default, it's `443` for `https` scheme and `80` for `http`                                                                        |  int   |    N     |
-| `roles`              | Mapping of catalog name to a role                                                                                                                                         |  dict  |    N     |
-| `http_headers`       | Additional HTTP headers to send with each request.                                                                                                                        |  dict  |    N     |
-| `session_properties` | Trino session properties. Run `SHOW SESSION` to see all options.                                                                                                          |  dict  |    N     |
-| `retries`            | Number of retries to attempt when a request fails. Default: `3`                                                                                                           |  int   |    N     |
-| `timezone`           | Timezone to use for the connection. Default: client-side local timezone                                                                                                   | string |    N     |
+| Option                    | Description                                                                                                                                                                             |  Type  | Required |
+|---------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------:|:--------:|
+| `type`                    | Engine type name - must be `trino`                                                                                                                                                      | string |    Y     |
+| `user`                    | The username (of the account) to log in to your cluster. When connecting to Starburst Galaxy clusters, you must include the role of the user as a suffix to the username.               | string |    Y     |
+| `host`                    | The hostname of your cluster. Don't include the `http://` or `https://` prefix.                                                                                                         | string |    Y     |
+| `catalog`                 | The name of a catalog in your cluster.                                                                                                                                                  | string |    Y     |
+| `http_scheme`             | The HTTP scheme to use when connecting to your cluster. By default, it's `https` and can only be `http` for no-auth or basic auth.                                                      | string |    N     |
+| `port`                    | The port to connect to your cluster. By default, it's `443` for `https` scheme and `80` for `http`                                                                                      |  int   |    N     |
+| `roles`                   | Mapping of catalog name to a role                                                                                                                                                       |  dict  |    N     |
+| `http_headers`            | Additional HTTP headers to send with each request.                                                                                                                                      |  dict  |    N     |
+| `session_properties`      | Trino session properties. Run `SHOW SESSION` to see all options.                                                                                                                        |  dict  |    N     |
+| `retries`                 | Number of retries to attempt when a request fails. Default: `3`                                                                                                                         |  int   |    N     |
+| `timezone`                | Timezone to use for the connection. Default: client-side local timezone                                                                                                                 | string |    N     |
+| `schema_location_mapping` | A mapping of regex patterns to S3 locations to use for the `LOCATION` property when creating schemas. See [Table and Schema locations](#table-and-schema-locations) for more details.   |  dict  |    N     |
+| `catalog_type_overrides`  | A mapping of catalog names to their connector type. This is used to enable/disable connector specific behavior. See [Catalog Type Overrides](#catalog-type-overrides) for more details. |  dict  |    N     |
 
 ## Table and Schema locations
 
@@ -204,6 +206,25 @@ SELECT ...
 
 This will cause SQLMesh to set the specified `LOCATION` when issuing a `CREATE TABLE` statement.
 
+## Catalog Type Overrides
+
+SQLMesh attempts to determine the connector type of a catalog by querying the `system.metadata.catalogs` table and checking the `connector_name` column.
+It checks if the connector name is `hive` for Hive connector behavior or contains `iceberg` or `delta_lake` for Iceberg or Delta Lake connector behavior respectively.
+However, the connector name may not always be a reliable way to determine the connector type, for example when using a custom connector or a fork of an existing connector.
+To handle such cases, you can use the `catalog_type_overrides` connection property to explicitly specify the connector type for specific catalogs.
+For example, to specify that the `datalake` catalog is using the Iceberg connector and the `analytics` catalog is using the Hive connector, you can configure the connection as follows:
+
+```yaml title="config.yaml"
+gateways:
+  trino:
+    connection:
+      type: trino
+      ...
+      catalog_type_overrides:
+        datalake: iceberg
+        analytics: hive
+```
+
 ## Authentication
 
 === "No Auth"

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "requests",
     "rich[jupyter]",
     "ruamel.yaml",
-    "sqlglot[rs]~=27.13.2",
+    "sqlglot[rs]~=27.15.3",
     "tenacity",
     "time-machine",
     "json-stream"
@@ -56,7 +56,7 @@ dev = [
     "agate",
     "beautifulsoup4",
     "clickhouse-connect",
-    "cryptography",
+    "cryptography<46.0.0",
     "databricks-sql-connector",
     "dbt-bigquery",
     "dbt-core",
@@ -119,7 +119,7 @@ postgres = ["psycopg2"]
 redshift = ["redshift_connector"]
 slack = ["slack_sdk"]
 snowflake = [
-    "cryptography",
+    "cryptography<46.0.0",
     "snowflake-connector-python[pandas,secure-local-storage]",
     "snowflake-snowpark-python",
 ]
@@ -135,7 +135,7 @@ lsp = [
     # Duplicate of web
     "fastapi==0.115.5",
     "watchfiles>=0.19.0",
-    "uvicorn[standard]==0.22.0",
+    # "uvicorn[standard]==0.22.0",
     "sse-starlette>=0.2.2",
     "pyarrow",
     # For lsp
@@ -225,7 +225,8 @@ module = [
     "pydantic_core.*",
     "dlt.*",
     "bigframes.*",
-    "json_stream.*"
+    "json_stream.*",
+    "duckdb.*"
 ]
 ignore_missing_imports = true
 

sqlmesh/core/config/connection.py

@@ -101,6 +101,7 @@ class ConnectionConfig(abc.ABC, BaseConfig):
     pre_ping: bool
     pretty_sql: bool = False
     schema_differ_overrides: t.Optional[t.Dict[str, t.Any]] = None
+    catalog_type_overrides: t.Optional[t.Dict[str, str]] = None
 
     # Whether to share a  single connection across threads or create a new connection per thread.
     shared_connection: t.ClassVar[bool] = False
@@ -176,6 +177,7 @@ def create_engine_adapter(
             pretty_sql=self.pretty_sql,
             shared_connection=self.shared_connection,
             schema_differ_overrides=self.schema_differ_overrides,
+            catalog_type_overrides=self.catalog_type_overrides,
             **self._extra_engine_config,
         )
 

sqlmesh/core/console.py

@@ -551,6 +551,22 @@ def log_skipped_models(self, snapshot_names: t.Set[str]) -> None:
     def log_failed_models(self, errors: t.List[NodeExecutionFailedError]) -> None:
         """Display list of models that failed during evaluation to the user."""
 
+    @abc.abstractmethod
+    def log_models_updated_during_restatement(
+        self,
+        snapshots: t.List[t.Tuple[SnapshotTableInfo, SnapshotTableInfo]],
+        environment_naming_info: EnvironmentNamingInfo,
+        default_catalog: t.Optional[str],
+    ) -> None:
+        """Display a list of models where new versions got deployed to the specified :environment while we were restating data the old versions
+
+        Args:
+            snapshots: a list of (snapshot_we_restated, snapshot_it_got_replaced_with_during_restatement) tuples
+            environment: which environment got updated while we were restating models
+            environment_naming_info: how snapshots are named in that :environment (for display name purposes)
+            default_catalog: the configured default catalog (for display name purposes)
+        """
+
     @abc.abstractmethod
     def loading_start(self, message: t.Optional[str] = None) -> uuid.UUID:
         """Starts loading and returns a unique ID that can be used to stop the loading. Optionally can display a message."""
@@ -771,6 +787,14 @@ def log_skipped_models(self, snapshot_names: t.Set[str]) -> None:
     def log_failed_models(self, errors: t.List[NodeExecutionFailedError]) -> None:
         pass
 
+    def log_models_updated_during_restatement(
+        self,
+        snapshots: t.List[t.Tuple[SnapshotTableInfo, SnapshotTableInfo]],
+        environment_naming_info: EnvironmentNamingInfo,
+        default_catalog: t.Optional[str],
+    ) -> None:
+        pass
+
     def log_destructive_change(
         self,
         snapshot_name: str,
@@ -1998,7 +2022,34 @@ def _prompt_categorize(
         plan = plan_builder.build()
 
         if plan.restatements:
-            self._print("\n[bold]Restating models\n")
+            # A plan can have restatements for the following reasons:
+            # - The user specifically called `sqlmesh plan` with --restate-model.
+            #   This creates a "restatement plan" which disallows all other changes and simply force-backfills
+            #   the selected models and their downstream dependencies using the versions of the models stored in state.
+            # - There are no specific restatements (so changes are allowed) AND dev previews need to be computed.
+            #   The "restatements" feature is currently reused for dev previews.
+            if plan.selected_models_to_restate:
+                # There were legitimate restatements, no dev previews
+                tree = Tree(
+                    "[bold]Models selected for restatement:[/bold]\n"
+                    "This causes backfill of the model itself as well as affected downstream models"
+                )
+                model_fqn_to_snapshot = {s.name: s for s in plan.snapshots.values()}
+                for model_fqn in plan.selected_models_to_restate:
+                    snapshot = model_fqn_to_snapshot[model_fqn]
+                    display_name = snapshot.display_name(
+                        plan.environment_naming_info,
+                        default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None,
+                        dialect=self.dialect,
+                    )
+                    tree.add(
+                        display_name
+                    )  # note: we deliberately dont show any intervals here; they get shown in the backfill section
+                self._print(tree)
+            else:
+                # We are computing dev previews, do not confuse the user by printing out something to do
+                # with restatements. Dev previews are already highlighted in the backfill step
+                pass
         else:
             self.show_environment_difference_summary(
                 plan.context_diff,
@@ -2225,6 +2276,30 @@ def log_failed_models(self, errors: t.List[NodeExecutionFailedError]) -> None:
             for node_name, msg in error_messages.items():
                 self._print(f"  [red]{node_name}[/red]\n\n{msg}")
 
+    def log_models_updated_during_restatement(
+        self,
+        snapshots: t.List[t.Tuple[SnapshotTableInfo, SnapshotTableInfo]],
+        environment_naming_info: EnvironmentNamingInfo,
+        default_catalog: t.Optional[str] = None,
+    ) -> None:
+        if snapshots:
+            tree = Tree(
+                f"[yellow]The following models had new versions deployed while data was being restated:[/yellow]"
+            )
+
+            for restated_snapshot, updated_snapshot in snapshots:
+                display_name = restated_snapshot.display_name(
+                    environment_naming_info,
+                    default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None,
+                    dialect=self.dialect,
+                )
+                current_branch = tree.add(display_name)
+                current_branch.add(f"restated version: '{restated_snapshot.version}'")
+                current_branch.add(f"currently active version: '{updated_snapshot.version}'")
+
+            self._print(tree)
+            self._print("")  # newline spacer
+
     def log_destructive_change(
         self,
         snapshot_name: str,

sqlmesh/core/context.py

@@ -1429,6 +1429,7 @@ def plan_builder(
         explain: t.Optional[bool] = None,
         ignore_cron: t.Optional[bool] = None,
         min_intervals: t.Optional[int] = None,
+        always_include_local_changes: t.Optional[bool] = None,
     ) -> PlanBuilder:
         """Creates a plan builder.
 
@@ -1467,6 +1468,8 @@ def plan_builder(
             diff_rendered: Whether the diff should compare raw vs rendered models
             min_intervals: Adjust the plan start date on a per-model basis in order to ensure at least this many intervals are covered
                 on every model when checking for missing intervals
+            always_include_local_changes: Usually when restatements are present, local changes in the filesystem are ignored.
+                However, it can be desirable to deploy changes + restatements in the same plan, so this flag overrides the default behaviour.
 
         Returns:
             The plan builder.
@@ -1583,13 +1586,20 @@ def plan_builder(
                 "Selector did not return any models. Please check your model selection and try again."
             )
 
+        if always_include_local_changes is None:
+            # default behaviour - if restatements are detected; we operate entirely out of state and ignore local changes
+            force_no_diff = restate_models is not None or (
+                backfill_models is not None and not backfill_models
+            )
+        else:
+            force_no_diff = not always_include_local_changes
+
         snapshots = self._snapshots(models_override)
         context_diff = self._context_diff(
             environment or c.PROD,
             snapshots=snapshots,
             create_from=create_from,
-            force_no_diff=restate_models is not None
-            or (backfill_models is not None and not backfill_models),
+            force_no_diff=force_no_diff,
             ensure_finalized_snapshots=self.config.plan.use_finalized_state,
             diff_rendered=diff_rendered,
             always_recreate_environment=self.config.plan.always_recreate_environment,
@@ -1644,13 +1654,22 @@ def plan_builder(
         elif forward_only is None:
             forward_only = self.config.plan.forward_only
 
+        # When handling prod restatements, only clear intervals from other model versions if we are using full virtual environments
+        # If we are not, then there is no point, because none of the data in dev environments can be promoted by definition
+        restate_all_snapshots = (
+            expanded_restate_models is not None
+            and not is_dev
+            and self.config.virtual_environment_mode.is_full
+        )
+
         return self.PLAN_BUILDER_TYPE(
             context_diff=context_diff,
             start=start,
             end=end,
             execution_time=execution_time,
             apply=self.apply,
             restate_models=expanded_restate_models,
+            restate_all_snapshots=restate_all_snapshots,
             backfill_models=backfill_models,
             no_gaps=no_gaps,
             skip_backfill=skip_backfill,

sqlmesh/core/dialect.py

@@ -174,6 +174,7 @@ def _parse_id_var(
 
     while (
         identifier
+        and not identifier.args.get("quoted")
         and self._is_connected()
         and (
             self._match_texts(("{", SQLMESH_MACRO_PREFIX))
@@ -349,13 +350,15 @@ def _parse_select(
     parse_subquery_alias: bool = True,
     parse_set_operation: bool = True,
     consume_pipe: bool = True,
+    from_: t.Optional[exp.From] = None,
 ) -> t.Optional[exp.Expression]:
     select = self.__parse_select(  # type: ignore
         nested=nested,
         table=table,
         parse_subquery_alias=parse_subquery_alias,
         parse_set_operation=parse_set_operation,
         consume_pipe=consume_pipe,
+        from_=from_,
     )
 
     if (