Merge branch 'main' into indirect_change_to_materialized_view_is_breaking

Author: xardasos

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"

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

@@ -2022,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,

sqlmesh/core/engine_adapter/base.py

@@ -223,6 +223,10 @@ def schema_differ(self) -> SchemaDiffer:
             }
         )
 
+    @property
+    def _catalog_type_overrides(self) -> t.Dict[str, str]:
+        return self._extra_config.get("catalog_type_overrides") or {}
+
     @classmethod
     def _casted_columns(
         cls,
@@ -430,7 +434,11 @@ def get_catalog_type(self, catalog: t.Optional[str]) -> str:
             raise UnsupportedCatalogOperationError(
                 f"{self.dialect} does not support catalogs and a catalog was provided: {catalog}"
             )
-        return self.DEFAULT_CATALOG_TYPE
+        return (
+            self._catalog_type_overrides.get(catalog, self.DEFAULT_CATALOG_TYPE)
+            if catalog
+            else self.DEFAULT_CATALOG_TYPE
+        )
 
     def get_catalog_type_from_table(self, table: TableName) -> str:
         """Get the catalog type from a table name if it has a catalog specified, otherwise return the current catalog type"""
@@ -1633,6 +1641,30 @@ def _insert_overwrite_by_condition(
                             target_columns_to_types=target_columns_to_types,
                             order_projections=False,
                         )
+                    elif insert_overwrite_strategy.is_merge:
+                        columns = [exp.column(col) for col in target_columns_to_types]
+                        when_not_matched_by_source = exp.When(
+                            matched=False,
+                            source=True,
+                            condition=where,
+                            then=exp.Delete(),
+                        )
+                        when_not_matched_by_target = exp.When(
+                            matched=False,
+                            source=False,
+                            then=exp.Insert(
+                                this=exp.Tuple(expressions=columns),
+                                expression=exp.Tuple(expressions=columns),
+                            ),
+                        )
+                        self._merge(
+                            target_table=table_name,
+                            query=query,
+                            on=exp.false(),
+                            whens=exp.Whens(
+                                expressions=[when_not_matched_by_source, when_not_matched_by_target]
+                            ),
+                        )
                     else:
                         insert_exp = exp.insert(
                             query,

sqlmesh/core/engine_adapter/bigquery.py

@@ -9,7 +9,6 @@
 
 from sqlmesh.core.dialect import to_schema
 from sqlmesh.core.engine_adapter.mixins import (
-    InsertOverwriteWithMergeMixin,
     ClusteredByMixin,
     RowDiffMixin,
     TableAlterClusterByOperation,
@@ -20,6 +19,7 @@
     DataObjectType,
     SourceQuery,
     set_catalog,
+    InsertOverwriteStrategy,
 )
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.schema_diff import TableAlterOperation, NestedSupport
@@ -54,7 +54,7 @@
 
 
 @set_catalog()
-class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, RowDiffMixin):
+class BigQueryEngineAdapter(ClusteredByMixin, RowDiffMixin):
     """
     BigQuery Engine Adapter using the `google-cloud-bigquery` library's DB API.
     """
@@ -68,6 +68,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
     MAX_COLUMN_COMMENT_LENGTH = 1024
     SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["SCHEMA"]
+    INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.MERGE
 
     SCHEMA_DIFFER_KWARGS = {
         "compatible_types": {

sqlmesh/core/engine_adapter/fabric.py

@@ -7,22 +7,15 @@
 from functools import cached_property
 from sqlglot import exp
 from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_result
+from sqlmesh.core.engine_adapter.mixins import LogicalMergeMixin
 from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter
 from sqlmesh.core.engine_adapter.shared import (
     InsertOverwriteStrategy,
-    SourceQuery,
 )
-from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.connection_pool import ConnectionPool
 
 
-if t.TYPE_CHECKING:
-    from sqlmesh.core._typing import TableName
-
-
-from sqlmesh.core.engine_adapter.mixins import LogicalMergeMixin
-
 logger = logging.getLogger(__name__)
 
 
@@ -58,26 +51,6 @@ def _target_catalog(self) -> t.Optional[str]:
     def _target_catalog(self, value: t.Optional[str]) -> None:
         self._connection_pool.set_attribute("target_catalog", value)
 
-    def _insert_overwrite_by_condition(
-        self,
-        table_name: TableName,
-        source_queries: t.List[SourceQuery],
-        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
-        where: t.Optional[exp.Condition] = None,
-        insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None,
-        **kwargs: t.Any,
-    ) -> None:
-        # Override to avoid MERGE statement which isn't fully supported in Fabric
-        return EngineAdapter._insert_overwrite_by_condition(
-            self,
-            table_name=table_name,
-            source_queries=source_queries,
-            target_columns_to_types=target_columns_to_types,
-            where=where,
-            insert_overwrite_strategy_override=InsertOverwriteStrategy.DELETE_INSERT,
-            **kwargs,
-        )
-
     @property
     def api_client(self) -> FabricHttpClient:
         # the requests Session is not guaranteed to be threadsafe

sqlmesh/core/engine_adapter/mixins.py

@@ -9,7 +9,6 @@
 from sqlglot.helper import seq_get
 
 from sqlmesh.core.engine_adapter.base import EngineAdapter
-from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy, SourceQuery
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.dialect import schema_
 from sqlmesh.core.schema_diff import TableAlterOperation
@@ -75,52 +74,6 @@ def _fetch_native_df(
         return df
 
 
-class InsertOverwriteWithMergeMixin(EngineAdapter):
-    def _insert_overwrite_by_condition(
-        self,
-        table_name: TableName,
-        source_queries: t.List[SourceQuery],
-        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
-        where: t.Optional[exp.Condition] = None,
-        insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None,
-        **kwargs: t.Any,
-    ) -> None:
-        """
-        Some engines do not support `INSERT OVERWRITE` but instead support
-        doing an "INSERT OVERWRITE" using a Merge expression but with the
-        predicate being `False`.
-        """
-        target_columns_to_types = target_columns_to_types or self.columns(table_name)
-        for source_query in source_queries:
-            with source_query as query:
-                query = self._order_projections_and_filter(
-                    query, target_columns_to_types, where=where
-                )
-                columns = [exp.column(col) for col in target_columns_to_types]
-                when_not_matched_by_source = exp.When(
-                    matched=False,
-                    source=True,
-                    condition=where,
-                    then=exp.Delete(),
-                )
-                when_not_matched_by_target = exp.When(
-                    matched=False,
-                    source=False,
-                    then=exp.Insert(
-                        this=exp.Tuple(expressions=columns),
-                        expression=exp.Tuple(expressions=columns),
-                    ),
-                )
-                self._merge(
-                    target_table=table_name,
-                    query=query,
-                    on=exp.false(),
-                    whens=exp.Whens(
-                        expressions=[when_not_matched_by_source, when_not_matched_by_target]
-                    ),
-                )
-
-
 class HiveMetastoreTablePropertiesMixin(EngineAdapter):
     MAX_TABLE_COMMENT_LENGTH = 4000
     MAX_COLUMN_COMMENT_LENGTH = 4000

sqlmesh/core/engine_adapter/mssql.py

@@ -16,7 +16,6 @@
 )
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
-    InsertOverwriteWithMergeMixin,
     PandasNativeFetchDFSupportMixin,
     VarcharSizeWorkaroundMixin,
     RowDiffMixin,
@@ -41,7 +40,6 @@
 class MSSQLEngineAdapter(
     EngineAdapterWithIndexSupport,
     PandasNativeFetchDFSupportMixin,
-    InsertOverwriteWithMergeMixin,
     GetCurrentCatalogFromFunctionMixin,
     VarcharSizeWorkaroundMixin,
     RowDiffMixin,
@@ -74,6 +72,7 @@ class MSSQLEngineAdapter(
         },
     }
     VARIABLE_LENGTH_DATA_TYPES = {"binary", "varbinary", "char", "varchar", "nchar", "nvarchar"}
+    INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.MERGE
 
     @property
     def catalog_support(self) -> CatalogSupport:

sqlmesh/core/engine_adapter/shared.py

@@ -243,6 +243,8 @@ class InsertOverwriteStrategy(Enum):
     # Issue a single INSERT query to replace a data range. The assumption is that the query engine will transparently match partition bounds
     # and replace data rather than append to it. Trino is an example of this when `hive.insert-existing-partitions-behavior=OVERWRITE` is configured
     INTO_IS_OVERWRITE = 4
+    # Do the INSERT OVERWRITE using merge since the engine doesn't support it natively
+    MERGE = 5
 
     @property
     def is_delete_insert(self) -> bool:
@@ -260,6 +262,10 @@ def is_replace_where(self) -> bool:
     def is_into_is_overwrite(self) -> bool:
         return self == InsertOverwriteStrategy.INTO_IS_OVERWRITE
 
+    @property
+    def is_merge(self) -> bool:
+        return self == InsertOverwriteStrategy.MERGE
+
 
 class SourceQuery:
     def __init__(