sqlmesh 0.228.1.dev1__py3-none-any.whl → 0.228.1.dev2__py3-none-any.whl

sqlmesh/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.228.1.dev1'
-__version_tuple__ = version_tuple = (0, 228, 1, 'dev1')
+__version__ = version = '0.228.1.dev2'
+__version_tuple__ = version_tuple = (0, 228, 1, 'dev2')
 
 __commit_id__ = commit_id = None

sqlmesh/core/config/connection.py

@@ -17,6 +17,7 @@ from pydantic_core import from_json
 from packaging import version
 from sqlglot import exp
 from sqlglot.helper import subclasses
+from sqlglot.errors import ParseError
 
 from sqlmesh.core import engine_adapter
 from sqlmesh.core.config.base import BaseConfig
@@ -1890,6 +1891,7 @@ class TrinoConnectionConfig(ConnectionConfig):
 
     # SQLMesh options
     schema_location_mapping: t.Optional[dict[re.Pattern, str]] = None
+    timestamp_mapping: t.Optional[dict[exp.DataType, exp.DataType]] = None
     concurrent_tasks: int = 4
     register_comments: bool = True
     pre_ping: t.Literal[False] = False
@@ -1914,6 +1916,34 @@ class TrinoConnectionConfig(ConnectionConfig):
                 )
         return compiled
 
+    @field_validator("timestamp_mapping", mode="before")
+    @classmethod
+    def _validate_timestamp_mapping(
+        cls, value: t.Optional[dict[str, str]]
+    ) -> t.Optional[dict[exp.DataType, exp.DataType]]:
+        if value is None:
+            return value
+
+        result: dict[exp.DataType, exp.DataType] = {}
+        for source_type, target_type in value.items():
+            try:
+                source_datatype = exp.DataType.build(source_type)
+            except ParseError:
+                raise ConfigError(
+                    f"Invalid SQL type string in timestamp_mapping: "
+                    f"'{source_type}' is not a valid SQL data type."
+                )
+            try:
+                target_datatype = exp.DataType.build(target_type)
+            except ParseError:
+                raise ConfigError(
+                    f"Invalid SQL type string in timestamp_mapping: "
+                    f"'{target_type}' is not a valid SQL data type."
+                )
+            result[source_datatype] = target_datatype
+
+        return result
+
     @model_validator(mode="after")
     def _root_validator(self) -> Self:
         port = self.port
@@ -2016,7 +2046,10 @@ class TrinoConnectionConfig(ConnectionConfig):
 
     @property
     def _extra_engine_config(self) -> t.Dict[str, t.Any]:
-        return {"schema_location_mapping": self.schema_location_mapping}
+        return {
+            "schema_location_mapping": self.schema_location_mapping,
+            "timestamp_mapping": self.timestamp_mapping,
+        }
 
 
 class ClickhouseConnectionConfig(ConnectionConfig):

sqlmesh/core/engine_adapter/trino.py

@@ -74,6 +74,32 @@ class TrinoEngineAdapter(
     def schema_location_mapping(self) -> t.Optional[t.Dict[re.Pattern, str]]:
         return self._extra_config.get("schema_location_mapping")
 
+    @property
+    def timestamp_mapping(self) -> t.Optional[t.Dict[exp.DataType, exp.DataType]]:
+        return self._extra_config.get("timestamp_mapping")
+
+    def _apply_timestamp_mapping(
+        self, columns_to_types: t.Dict[str, exp.DataType]
+    ) -> t.Tuple[t.Dict[str, exp.DataType], t.Set[str]]:
+        """Apply custom timestamp mapping to column types.
+
+        Returns:
+            A tuple of (mapped_columns_to_types, mapped_column_names) where mapped_column_names
+            contains the names of columns that were found in the mapping.
+        """
+        if not self.timestamp_mapping:
+            return columns_to_types, set()
+
+        result = {}
+        mapped_columns: t.Set[str] = set()
+        for column, column_type in columns_to_types.items():
+            if column_type in self.timestamp_mapping:
+                result[column] = self.timestamp_mapping[column_type]
+                mapped_columns.add(column)
+            else:
+                result[column] = column_type
+        return result, mapped_columns
+
     @property
     def catalog_support(self) -> CatalogSupport:
         return CatalogSupport.FULL_SUPPORT
@@ -117,7 +143,7 @@ class TrinoEngineAdapter(
         try:
             yield
         finally:
-            self.execute(f"RESET SESSION AUTHORIZATION")
+            self.execute("RESET SESSION AUTHORIZATION")
 
     def replace_query(
         self,
@@ -286,8 +312,11 @@ class TrinoEngineAdapter(
         is_view: bool = False,
         materialized: bool = False,
     ) -> exp.Schema:
+        target_columns_to_types, mapped_columns = self._apply_timestamp_mapping(
+            target_columns_to_types
+        )
         if "delta_lake" in self.get_catalog_type_from_table(table):
-            target_columns_to_types = self._to_delta_ts(target_columns_to_types)
+            target_columns_to_types = self._to_delta_ts(target_columns_to_types, mapped_columns)
 
         return super()._build_schema_exp(
             table, target_columns_to_types, column_descriptions, expressions, is_view
@@ -313,10 +342,15 @@ class TrinoEngineAdapter(
         source_columns: t.Optional[t.List[str]] = None,
         **kwargs: t.Any,
     ) -> None:
+        mapped_columns: t.Set[str] = set()
+        if target_columns_to_types:
+            target_columns_to_types, mapped_columns = self._apply_timestamp_mapping(
+                target_columns_to_types
+            )
         if target_columns_to_types and "delta_lake" in self.get_catalog_type_from_table(
             target_table
         ):
-            target_columns_to_types = self._to_delta_ts(target_columns_to_types)
+            target_columns_to_types = self._to_delta_ts(target_columns_to_types, mapped_columns)
 
         return super()._scd_type_2(
             target_table,
@@ -346,18 +380,21 @@ class TrinoEngineAdapter(
     # - `timestamp(3) with time zone` for timezone-aware
     # https://trino.io/docs/current/connector/delta-lake.html#delta-lake-to-trino-type-mapping
     def _to_delta_ts(
-        self, columns_to_types: t.Dict[str, exp.DataType]
+        self,
+        columns_to_types: t.Dict[str, exp.DataType],
+        skip_columns: t.Optional[t.Set[str]] = None,
     ) -> t.Dict[str, exp.DataType]:
         ts6 = exp.DataType.build("timestamp(6)")
         ts3_tz = exp.DataType.build("timestamp(3) with time zone")
+        skip = skip_columns or set()
 
         delta_columns_to_types = {
-            k: ts6 if v.is_type(exp.DataType.Type.TIMESTAMP) else v
+            k: ts6 if k not in skip and v.is_type(exp.DataType.Type.TIMESTAMP) else v
             for k, v in columns_to_types.items()
         }
 
         delta_columns_to_types = {
-            k: ts3_tz if v.is_type(exp.DataType.Type.TIMESTAMPTZ) else v
+            k: ts3_tz if k not in skip and v.is_type(exp.DataType.Type.TIMESTAMPTZ) else v
             for k, v in delta_columns_to_types.items()
         }
 

{sqlmesh-0.228.1.dev1.dist-info → sqlmesh-0.228.1.dev2.dist-info}/METADATA

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

{sqlmesh-0.228.1.dev1.dist-info → sqlmesh-0.228.1.dev2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 sqlmesh/__init__.py,sha256=v_spqQEhcnGaahp1yPvMqUIa6mhH3cs3Bc1CznxvCEA,7965
-sqlmesh/_version.py,sha256=G0EnPfaObCiXxYjW9mQvEf7Bpf-ACOLqs4xEIF8Z_hg,721
+sqlmesh/_version.py,sha256=-UbxvaOsHNt41ClgzojQDAjyZAdfbkbOj_p0-3b_vJw,721
 sqlmesh/magics.py,sha256=7Q1_lXSD_PgYH40Hsx6-OkfSQC3UJZgF043RVFRnw1s,42082
 sqlmesh/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlmesh/cicd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -42,7 +42,7 @@ sqlmesh/core/config/__init__.py,sha256=tnEakbd8FAgSLYmjzuYAAgHIpJ00lwMKAhD_Cfs2O
 sqlmesh/core/config/base.py,sha256=t8NQmsgQoZSc-k0dlDiCb8t1jj0AMYdGZ-6se9q_Pks,4898
 sqlmesh/core/config/categorizer.py,sha256=6vzUoNLjR6GOEb_2mYVz2TwmMv2BfldgHX2u-Le5HZs,1975
 sqlmesh/core/config/common.py,sha256=9V6PltBAjYeWLOU5dAbqL55BSFfpg8z8t2Op1x_PLhU,6418
-sqlmesh/core/config/connection.py,sha256=qynbsjRT_KOI6InKdCIXdqGXExF43EuawQp0NxTHPG4,91168
+sqlmesh/core/config/connection.py,sha256=l2GUpZtCJyVk94JiTcvZbrrH9dpOfZChCwyM6Z20Efs,92425
 sqlmesh/core/config/dbt.py,sha256=xSQ4NEVWhZj_aRYpyy4MWcRJ8Qa0o28w2ZBLI4bs3_I,468
 sqlmesh/core/config/format.py,sha256=6CXFbvnor56xbldKE-Vrm9k_ABRoY4v6vgIb3mCihiQ,1355
 sqlmesh/core/config/gateway.py,sha256=tYngyqwd_4Qr9lhcv2hlvLvb_2pgYYtKu6hdGsTr-4I,1931
@@ -76,7 +76,7 @@ sqlmesh/core/engine_adapter/risingwave.py,sha256=d_1MxpXNONyyLnuELa7bILkJlLquf4j
 sqlmesh/core/engine_adapter/shared.py,sha256=bM4GJSAR0dU3wCqsTl2SIcy2j_8BGusQvnme99l6wnE,13701
 sqlmesh/core/engine_adapter/snowflake.py,sha256=6rMuhuhp2K-UH8dVnmiieucfOevxmK8vR3N5-dj4MDA,33453
 sqlmesh/core/engine_adapter/spark.py,sha256=ZDEg4rx_cvPcLG83PSWu5nkXzChaCbmb7ka2J2ngEEU,23068
-sqlmesh/core/engine_adapter/trino.py,sha256=dy5DXCbhPB9nuTDytiNuDXe-wGgxv5b9CjhEVpbQCq4,18180
+sqlmesh/core/engine_adapter/trino.py,sha256=F6Cs0RxbMnk5tOYxp3GdMNhK3nnBrtnU8gldMD3POh0,19731
 sqlmesh/core/linter/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlmesh/core/linter/definition.py,sha256=1EOhKdF16jmeqISfcrR-8fzMdgXuxpB7wb3QaepBPeU,5564
 sqlmesh/core/linter/helpers.py,sha256=cwKXP4sL6azRtNVGbMfJ5_6Hqq5Xx2M2rRLCgH3Y3ag,10743
@@ -238,7 +238,7 @@ sqlmesh/utils/pydantic.py,sha256=-yppkVlw6iSBaSiKjbe7OChxL-u3urOS4-KCjJEgsRU,120
 sqlmesh/utils/rich.py,sha256=cwQ5nJ6sgz64xHtoh6_ec7ReV5YpsOGhMtUJnwoRfEI,3549
 sqlmesh/utils/windows.py,sha256=0F9RdpuuCoG5NiEDXvWlAGCiJ-59OjSAmgFF5wW05aY,1133
 sqlmesh/utils/yaml.py,sha256=KFBd7hsKNRTtRudGR7d410qUYffQv0EWRcDM8hVNNZg,3025
-sqlmesh-0.228.1.dev1.dist-info/licenses/LICENSE,sha256=OlMefUjgWJdULtf84BLW0AZZcY8DwdgQqb_1j2862j8,11346
+sqlmesh-0.228.1.dev2.dist-info/licenses/LICENSE,sha256=OlMefUjgWJdULtf84BLW0AZZcY8DwdgQqb_1j2862j8,11346
 sqlmesh_dbt/__init__.py,sha256=awYS5y5mz-1NUmx6i5h5NSTJ7tidRl9NC0FAnFWSF6U,350
 sqlmesh_dbt/cli.py,sha256=p9foHjAW9ni7BTOJ2loynk47M0Sf43QIJZRggOzF5tc,6351
 sqlmesh_dbt/console.py,sha256=RwWLYnEZHzn9Xp-e2gbZvkdKbWbBLN146geI84mJitg,1132
@@ -363,8 +363,8 @@ web/server/api/endpoints/models.py,sha256=kwj0s7uve3iZSMfmjkoPVMFMeY1sD0peTeyrWf
 web/server/api/endpoints/modules.py,sha256=8hqqgonGay_mJmpCw0IdbjsPhWlQH2VLdKAqha-myac,468
 web/server/api/endpoints/plan.py,sha256=bbbY50W_2MsZSTxOHWMKz0tbIm75nsRSlPy8GI2fg9Q,9306
 web/server/api/endpoints/table_diff.py,sha256=8XTwgOh6QBbNy_hTM1JuHgRjbnie-pGPrphiW-FNLjQ,6058
-sqlmesh-0.228.1.dev1.dist-info/METADATA,sha256=62st3YlLnXg9N1lEeRS4cqt0ctsYq4Qv8LMHbBmq9q8,26685
-sqlmesh-0.228.1.dev1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sqlmesh-0.228.1.dev1.dist-info/entry_points.txt,sha256=sHAf6tQczIM8xZoduN4qaUjV7QEPVUUW_LCT8EDUMv4,155
-sqlmesh-0.228.1.dev1.dist-info/top_level.txt,sha256=RQ-33FPe2IgL0rgossAfJkCRtqslz9b7wFARqiWLC5Q,24
-sqlmesh-0.228.1.dev1.dist-info/RECORD,,
+sqlmesh-0.228.1.dev2.dist-info/METADATA,sha256=WTvsWqQryymDjFJv6vkk_r_XuidAGuO8X6K5guIFP4A,26685
+sqlmesh-0.228.1.dev2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sqlmesh-0.228.1.dev2.dist-info/entry_points.txt,sha256=sHAf6tQczIM8xZoduN4qaUjV7QEPVUUW_LCT8EDUMv4,155
+sqlmesh-0.228.1.dev2.dist-info/top_level.txt,sha256=RQ-33FPe2IgL0rgossAfJkCRtqslz9b7wFARqiWLC5Q,24
+sqlmesh-0.228.1.dev2.dist-info/RECORD,,