sqlmesh.core.engine_adapter.trino

  1from __future__ import annotations
  2
  3import contextlib
  4import re
  5import typing as t
  6from functools import lru_cache
  7
  8from sqlglot import exp
  9from sqlglot.helper import seq_get
 10from tenacity import retry, wait_fixed, stop_after_attempt, retry_if_result
 11
 12from sqlmesh.core.dialect import schema_, to_schema
 13from sqlmesh.core.engine_adapter.mixins import (
 14    GetCurrentCatalogFromFunctionMixin,
 15    HiveMetastoreTablePropertiesMixin,
 16    PandasNativeFetchDFSupportMixin,
 17    RowDiffMixin,
 18)
 19from sqlmesh.core.engine_adapter.shared import (
 20    CatalogSupport,
 21    CommentCreationTable,
 22    CommentCreationView,
 23    DataObject,
 24    DataObjectType,
 25    InsertOverwriteStrategy,
 26    SourceQuery,
 27    set_catalog,
 28)
 29from sqlmesh.utils import get_source_columns_to_types
 30from sqlmesh.utils.errors import SQLMeshError
 31from sqlmesh.utils.date import TimeLike
 32
 33if t.TYPE_CHECKING:
 34    from sqlmesh.core._typing import SchemaName, SessionProperties, TableName
 35    from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF
 36
 37CATALOG_TYPES_SUPPORTING_REPLACE_TABLE = {"iceberg", "delta_lake"}
 38
 39
 40@set_catalog()
 41class TrinoEngineAdapter(
 42    PandasNativeFetchDFSupportMixin,
 43    HiveMetastoreTablePropertiesMixin,
 44    GetCurrentCatalogFromFunctionMixin,
 45    RowDiffMixin,
 46):
 47    DIALECT = "trino"
 48    INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.INTO_IS_OVERWRITE
 49    # Trino does technically support transactions but it doesn't work correctly with partition overwrite so we
 50    # disable transactions. If we need to get them enabled again then we would need to disable auto commit on the
 51    # connector and then figure out how to get insert/overwrite to work correctly without it.
 52    SUPPORTS_TRANSACTIONS = False
 53    CURRENT_CATALOG_EXPRESSION = exp.column("current_catalog")
 54    COMMENT_CREATION_TABLE = CommentCreationTable.IN_SCHEMA_DEF_NO_CTAS
 55    COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
 56    SUPPORTS_REPLACE_TABLE = False
 57    SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["SCHEMA"]
 58    DEFAULT_CATALOG_TYPE = "hive"
 59    QUOTE_IDENTIFIERS_IN_VIEWS = False
 60    SUPPORTS_QUERY_EXECUTION_TRACKING = True
 61    SCHEMA_DIFFER_KWARGS = {
 62        "parameterized_type_defaults": {
 63            # default decimal precision varies across backends
 64            exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(), (0,)],
 65            exp.DataType.build("CHAR", dialect=DIALECT).this: [(1,)],
 66            exp.DataType.build("TIMESTAMP", dialect=DIALECT).this: [(3,)],
 67        },
 68    }
 69    # some catalogs support microsecond (precision 6) but it has to be specifically enabled (Hive) or just isnt available (Delta / TIMESTAMP WITH TIME ZONE)
 70    # and even if you have a TIMESTAMP(6) the date formatting functions still only support millisecond precision
 71    MAX_TIMESTAMP_PRECISION = 3
 72
 73    @property
 74    def schema_location_mapping(self) -> t.Optional[t.Dict[re.Pattern, str]]:
 75        return self._extra_config.get("schema_location_mapping")
 76
 77    @property
 78    def timestamp_mapping(self) -> t.Optional[t.Dict[exp.DataType, exp.DataType]]:
 79        return self._extra_config.get("timestamp_mapping")
 80
 81    def _apply_timestamp_mapping(
 82        self, columns_to_types: t.Dict[str, exp.DataType]
 83    ) -> t.Tuple[t.Dict[str, exp.DataType], t.Set[str]]:
 84        """Apply custom timestamp mapping to column types.
 85
 86        Returns:
 87            A tuple of (mapped_columns_to_types, mapped_column_names) where mapped_column_names
 88            contains the names of columns that were found in the mapping.
 89        """
 90        if not self.timestamp_mapping:
 91            return columns_to_types, set()
 92
 93        result = {}
 94        mapped_columns: t.Set[str] = set()
 95        for column, column_type in columns_to_types.items():
 96            if column_type in self.timestamp_mapping:
 97                result[column] = self.timestamp_mapping[column_type]
 98                mapped_columns.add(column)
 99            else:
100                result[column] = column_type
101        return result, mapped_columns
102
103    @property
104    def catalog_support(self) -> CatalogSupport:
105        return CatalogSupport.FULL_SUPPORT
106
107    def set_current_catalog(self, catalog: str) -> None:
108        """Sets the catalog name of the current connection."""
109        self.execute(exp.Use(this=schema_(db="information_schema", catalog=catalog)))
110
111    @lru_cache()
112    def get_catalog_type(self, catalog: t.Optional[str]) -> str:
113        row: t.Tuple = tuple()
114        if catalog:
115            if catalog_type_override := self._catalog_type_overrides.get(catalog):
116                return catalog_type_override
117            row = (
118                self.fetchone(
119                    f"select connector_name from system.metadata.catalogs where catalog_name='{catalog}'"
120                )
121                or ()
122            )
123        return seq_get(row, 0) or self.DEFAULT_CATALOG_TYPE
124
125    @contextlib.contextmanager
126    def session(self, properties: SessionProperties) -> t.Iterator[None]:
127        authorization = properties.get("authorization")
128        if not authorization:
129            yield
130            return
131
132        if not isinstance(authorization, exp.Expr):
133            authorization = exp.Literal.string(authorization)
134
135        if not authorization.is_string:
136            raise SQLMeshError(
137                "Invalid value for `session_properties.authorization`. Must be a string literal."
138            )
139
140        authorization_sql = authorization.sql(dialect=self.dialect)
141
142        self.execute(f"SET SESSION AUTHORIZATION {authorization_sql}")
143        try:
144            yield
145        finally:
146            self.execute("RESET SESSION AUTHORIZATION")
147
148    def replace_query(
149        self,
150        table_name: TableName,
151        query_or_df: QueryOrDF,
152        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
153        table_description: t.Optional[str] = None,
154        column_descriptions: t.Optional[t.Dict[str, str]] = None,
155        source_columns: t.Optional[t.List[str]] = None,
156        supports_replace_table_override: t.Optional[bool] = None,
157        **kwargs: t.Any,
158    ) -> None:
159        catalog_type = self.get_catalog_type_from_table(table_name)
160        # User may have a custom catalog type name so we are assuming they keep the catalog type still in the name
161        # Ex: `acme_iceberg` would be identified as an iceberg catalog and therefore supports replace table
162        supports_replace_table_override = None
163        for replace_table_catalog_type in CATALOG_TYPES_SUPPORTING_REPLACE_TABLE:
164            if replace_table_catalog_type in catalog_type:
165                supports_replace_table_override = True
166                break
167
168        super().replace_query(
169            table_name=table_name,
170            query_or_df=query_or_df,
171            target_columns_to_types=target_columns_to_types,
172            table_description=table_description,
173            column_descriptions=column_descriptions,
174            source_columns=source_columns,
175            supports_replace_table_override=supports_replace_table_override,
176            **kwargs,
177        )
178
179    def _insert_overwrite_by_condition(
180        self,
181        table_name: TableName,
182        source_queries: t.List[SourceQuery],
183        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
184        where: t.Optional[exp.Condition] = None,
185        insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None,
186        **kwargs: t.Any,
187    ) -> None:
188        catalog = exp.to_table(table_name).catalog or self.get_current_catalog()
189
190        if where and self.get_catalog_type(catalog) == "hive":
191            # These session properties are only valid for the Trino Hive connector
192            # Attempting to set them on an Iceberg catalog will throw an error:
193            # "Session property 'catalog.insert_existing_partitions_behavior' does not exist"
194            self.execute(f"SET SESSION {catalog}.insert_existing_partitions_behavior='OVERWRITE'")
195            super()._insert_overwrite_by_condition(
196                table_name, source_queries, target_columns_to_types, where
197            )
198            self.execute(f"SET SESSION {catalog}.insert_existing_partitions_behavior='APPEND'")
199        else:
200            super()._insert_overwrite_by_condition(
201                table_name,
202                source_queries,
203                target_columns_to_types,
204                where,
205                insert_overwrite_strategy_override=InsertOverwriteStrategy.DELETE_INSERT,
206            )
207
208    def _truncate_table(self, table_name: TableName) -> None:
209        table = exp.to_table(table_name)
210        # Some trino connectors don't support truncate so we use delete.
211        self.execute(f"DELETE FROM {table.sql(dialect=self.dialect, identify=True)}")
212
213    def _get_data_objects(
214        self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
215    ) -> t.List[DataObject]:
216        """
217        Returns all the data objects that exist in the given schema and optionally catalog.
218        """
219        schema_name = to_schema(schema_name)
220        schema = schema_name.db
221        catalog = schema_name.catalog or self.get_current_catalog()
222        query = (
223            exp.select(
224                exp.column("table_catalog", table="t").as_("catalog"),
225                exp.column("table_schema", table="t").as_("schema"),
226                exp.column("table_name", table="t").as_("name"),
227                exp.case()
228                .when(
229                    exp.column("name", table="mv").is_(exp.null()).not_(),
230                    exp.Literal.string("materialized_view"),
231                )
232                .when(
233                    exp.column("table_type", table="t").eq("BASE TABLE"),
234                    exp.Literal.string("table"),
235                )
236                .else_(exp.column("table_type", table="t"))
237                .as_("type"),
238            )
239            .from_(exp.to_table(f"{catalog}.information_schema.tables", alias="t"))
240            .join(
241                exp.to_table("system.metadata.materialized_views", alias="mv"),
242                on=exp.and_(
243                    exp.column("catalog_name", table="mv").eq(
244                        exp.column("table_catalog", table="t")
245                    ),
246                    exp.column("schema_name", table="mv").eq(exp.column("table_schema", table="t")),
247                    exp.column("name", table="mv").eq(exp.column("table_name", table="t")),
248                ),
249                join_type="left",
250            )
251            .where(
252                exp.and_(
253                    exp.column("table_schema", table="t").eq(schema),
254                    exp.or_(
255                        exp.column("catalog_name", table="mv").is_(exp.null()),
256                        exp.column("catalog_name", table="mv").eq(catalog),
257                    ),
258                    exp.or_(
259                        exp.column("schema_name", table="mv").is_(exp.null()),
260                        exp.column("schema_name", table="mv").eq(schema),
261                    ),
262                )
263            )
264        )
265        if object_names:
266            query = query.where(exp.column("table_name", table="t").isin(*object_names))
267        df = self.fetchdf(query)
268        return [
269            DataObject(
270                catalog=row.catalog,  # type: ignore
271                schema=row.schema,  # type: ignore
272                name=row.name,  # type: ignore
273                type=DataObjectType.from_str(row.type),  # type: ignore
274            )
275            for row in df.itertuples()
276        ]
277
278    def _df_to_source_queries(
279        self,
280        df: DF,
281        target_columns_to_types: t.Dict[str, exp.DataType],
282        batch_size: int,
283        target_table: TableName,
284        source_columns: t.Optional[t.List[str]] = None,
285    ) -> t.List[SourceQuery]:
286        import pandas as pd
287        from pandas.api.types import is_datetime64_any_dtype  # type: ignore
288
289        assert isinstance(df, pd.DataFrame)
290        source_columns_to_types = get_source_columns_to_types(
291            target_columns_to_types, source_columns
292        )
293
294        # Trino does not accept timestamps in ISOFORMAT that include the "T". `execution_time` is stored in
295        # Pandas with that format, so we convert the column to a string with the proper format and CAST to
296        # timestamp in Trino.
297        for column, kind in source_columns_to_types.items():
298            dtype = df.dtypes[column]
299            if is_datetime64_any_dtype(dtype) and getattr(dtype, "tz", None) is not None:
300                df[column] = pd.to_datetime(df[column]).map(lambda x: x.isoformat(" "))
301
302        return super()._df_to_source_queries(
303            df, target_columns_to_types, batch_size, target_table, source_columns=source_columns
304        )
305
306    def _build_schema_exp(
307        self,
308        table: exp.Table,
309        target_columns_to_types: t.Dict[str, exp.DataType],
310        column_descriptions: t.Optional[t.Dict[str, str]] = None,
311        expressions: t.Optional[t.List[exp.PrimaryKey]] = None,
312        is_view: bool = False,
313        materialized: bool = False,
314    ) -> exp.Schema:
315        target_columns_to_types, mapped_columns = self._apply_timestamp_mapping(
316            target_columns_to_types
317        )
318        if "delta_lake" in self.get_catalog_type_from_table(table):
319            target_columns_to_types = self._to_delta_ts(target_columns_to_types, mapped_columns)
320
321        return super()._build_schema_exp(
322            table, target_columns_to_types, column_descriptions, expressions, is_view
323        )
324
325    def _scd_type_2(
326        self,
327        target_table: TableName,
328        source_table: QueryOrDF,
329        unique_key: t.Sequence[exp.Expr],
330        valid_from_col: exp.Column,
331        valid_to_col: exp.Column,
332        execution_time: t.Union[TimeLike, exp.Column],
333        invalidate_hard_deletes: bool = True,
334        updated_at_col: t.Optional[exp.Column] = None,
335        check_columns: t.Optional[t.Union[exp.Star, t.Sequence[exp.Expr]]] = None,
336        updated_at_as_valid_from: bool = False,
337        execution_time_as_valid_from: bool = False,
338        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
339        table_description: t.Optional[str] = None,
340        column_descriptions: t.Optional[t.Dict[str, str]] = None,
341        truncate: bool = False,
342        source_columns: t.Optional[t.List[str]] = None,
343        **kwargs: t.Any,
344    ) -> None:
345        mapped_columns: t.Set[str] = set()
346        if target_columns_to_types:
347            target_columns_to_types, mapped_columns = self._apply_timestamp_mapping(
348                target_columns_to_types
349            )
350        if target_columns_to_types and "delta_lake" in self.get_catalog_type_from_table(
351            target_table
352        ):
353            target_columns_to_types = self._to_delta_ts(target_columns_to_types, mapped_columns)
354
355        return super()._scd_type_2(
356            target_table,
357            source_table,
358            unique_key,
359            valid_from_col,
360            valid_to_col,
361            execution_time,
362            invalidate_hard_deletes,
363            updated_at_col,
364            check_columns,
365            updated_at_as_valid_from,
366            execution_time_as_valid_from,
367            target_columns_to_types,
368            table_description,
369            column_descriptions,
370            truncate,
371            source_columns,
372            **kwargs,
373        )
374
375    # delta_lake only supports two timestamp data types. This method converts other
376    # timestamp types to those two for use in DDL statements. Trino/delta automatically
377    # converts the data values to the correct type on write, so we only need to handle
378    # the column types in DDL.
379    # - `timestamp(6)` for non-timezone-aware
380    # - `timestamp(3) with time zone` for timezone-aware
381    # https://trino.io/docs/current/connector/delta-lake.html#delta-lake-to-trino-type-mapping
382    def _to_delta_ts(
383        self,
384        columns_to_types: t.Dict[str, exp.DataType],
385        skip_columns: t.Optional[t.Set[str]] = None,
386    ) -> t.Dict[str, exp.DataType]:
387        ts6 = exp.DataType.build("timestamp(6)")
388        ts3_tz = exp.DataType.build("timestamp(3) with time zone")
389        skip = skip_columns or set()
390
391        delta_columns_to_types = {
392            k: ts6 if k not in skip and v.is_type(exp.DataType.Type.TIMESTAMP) else v
393            for k, v in columns_to_types.items()
394        }
395
396        delta_columns_to_types = {
397            k: ts3_tz if k not in skip and v.is_type(exp.DataType.Type.TIMESTAMPTZ) else v
398            for k, v in delta_columns_to_types.items()
399        }
400
401        return delta_columns_to_types
402
403    @retry(wait=wait_fixed(1), stop=stop_after_attempt(10), retry=retry_if_result(lambda v: not v))
404    def _block_until_table_exists(self, table_name: TableName) -> bool:
405        return self.table_exists(table_name)
406
407    def _create_schema(
408        self,
409        schema_name: SchemaName,
410        ignore_if_exists: bool,
411        warn_on_error: bool,
412        properties: t.List[exp.Expr],
413        kind: str,
414    ) -> None:
415        if mapped_location := self._schema_location(schema_name):
416            properties.append(exp.LocationProperty(this=exp.Literal.string(mapped_location)))
417
418        return super()._create_schema(
419            schema_name=schema_name,
420            ignore_if_exists=ignore_if_exists,
421            warn_on_error=warn_on_error,
422            properties=properties,
423            kind=kind,
424        )
425
426    def _create_table(
427        self,
428        table_name_or_schema: t.Union[exp.Schema, TableName],
429        expression: t.Optional[exp.Expr],
430        exists: bool = True,
431        replace: bool = False,
432        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
433        table_description: t.Optional[str] = None,
434        column_descriptions: t.Optional[t.Dict[str, str]] = None,
435        table_kind: t.Optional[str] = None,
436        track_rows_processed: bool = True,
437        **kwargs: t.Any,
438    ) -> None:
439        super()._create_table(
440            table_name_or_schema=table_name_or_schema,
441            expression=expression,
442            exists=exists,
443            replace=replace,
444            target_columns_to_types=target_columns_to_types,
445            table_description=table_description,
446            column_descriptions=column_descriptions,
447            table_kind=table_kind,
448            track_rows_processed=track_rows_processed,
449            **kwargs,
450        )
451
452        # extract the table name
453        if isinstance(table_name_or_schema, exp.Schema):
454            table_name = table_name_or_schema.this
455            assert isinstance(table_name, exp.Table)
456        else:
457            table_name = table_name_or_schema
458
459        if "hive" in self.get_catalog_type_from_table(table_name):
460            # the Trino Hive connector can take a few seconds for metadata changes to propagate to all internal threads
461            # (even if metadata TTL is set to 0s)
462            # Blocking until the table shows up means that subsequent code expecting it to exist immediately will not fail
463            self._block_until_table_exists(table_name)
464
465    def _schema_location(self, schema_name: SchemaName) -> t.Optional[str]:
466        if mapping := self.schema_location_mapping:
467            schema = to_schema(schema_name)
468            match_key = schema.db
469
470            # only consider the catalog if it is present
471            if schema.catalog:
472                match_key = f"{schema.catalog}.{match_key}"
473
474            for k, v in mapping.items():
475                if re.match(k, match_key):
476                    return v.replace("@{schema_name}", schema.db).replace(
477                        "@{catalog_name}", schema.catalog
478                    )
479        return None