Repository: TobikoData/sqlmesh
Pull request #2330 (2024-03-22 14:59:10)

Feat: support Trino Delta Lake connector
- Only tested with Hive catalog type
- Converts timestamp types in DDL schemas
    - delta_lake only supports two timestamp data types. We convert other timestamp types to those two for use in DDL statements. Trino/delta automatically converts the data values to the correct type on write, so we only need to handle the column types in DDL. https://trino.io/docs/current/connector/delta-lake.html#delta-lake-to-trino-type-mapping
        - `timestamp(6)` for non-timezone-aware
        - `timestamp(3) with time zone` for timezone-aware


Patch:
diff --git a/Makefile b/Makefile
index 069c2d84..bc921bcc 100644
--- a/Makefile
+++ b/Makefile
@@ -179,4 +179,4 @@ spark-pyspark-test:
 	pytest -n auto -m "spark_pyspark"
 
 trino-test:
-	pytest -n auto -m "trino or trino_iceberg"
\ No newline at end of file
+	pytest -n auto -m "trino or trino_iceberg or trino_delta"
\ No newline at end of file
diff --git a/docs/integrations/engines/trino.md b/docs/integrations/engines/trino.md
index 61ec1ae5..db1c16ae 100644
--- a/docs/integrations/engines/trino.md
+++ b/docs/integrations/engines/trino.md
@@ -17,7 +17,7 @@ pip install "trino[external-authentication-token-cache]"
 
 ### Trino Connector Support
 
-The trino engine adapter has been tested against the [Hive Connector](https://trino.io/docs/current/connector/hive.html) and the [Iceberg Connector](https://trino.io/docs/current/connector/iceberg.html).
+The trino engine adapter has been tested against the [Hive Connector](https://trino.io/docs/current/connector/hive.html), [Iceberg Connector](https://trino.io/docs/current/connector/iceberg.html), and [Delta Lake Connector](https://trino.io/docs/current/connector/delta-lake.html).
 
 Please let us know on [Slack](https://tobikodata.com/slack) if you are wanting to use another connector or have tried another connector.
 
@@ -49,6 +49,17 @@ iceberg.catalog.type=hive_metastore
 
 The `jdbc`, `rest` and `nessie` catalogs do not support views and are thus incompatible with SQLMesh.
 
+#### Delta Lake Connector Configuration
+
+The Trino adapter Delta Lake connector has only been tested with the Hive metastore catalog type.
+
+The [properties file](https://trino.io/docs/current/connector/delta-lake.html#general-configuration) must include the Hive metastore URI and catalog name in addition to any other [general properties](https://trino.io/docs/current/object-storage/metastores.html#general-metastore-properties).
+
+``` properties linenums="1"
+hive.metastore.uri=thrift://example.net:9083
+delta.hive-catalog-name=datalake_delta # example catalog name, can be any valid string
+```
+
 ### Connection options
 
 | Option               | Description                                                                                                                                                               |  Type  | Required |
diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py
index c425e1c1..2d8ecd77 100644
--- a/sqlmesh/core/engine_adapter/mixins.py
+++ b/sqlmesh/core/engine_adapter/mixins.py
@@ -224,8 +224,8 @@ class HiveMetastoreTablePropertiesMixin(EngineAdapter):
     def _truncate_comment(
         self, comment: str, length: t.Optional[int], escape_backslash: bool = False
     ) -> str:
-        # iceberg does not have a comment length limit
-        if self.current_catalog_type == "iceberg":
+        # iceberg and delta do not have a comment length limit
+        if self.current_catalog_type in ("iceberg", "delta"):
             return comment
         return super()._truncate_comment(comment, length, escape_backslash)
 
diff --git a/sqlmesh/core/engine_adapter/trino.py b/sqlmesh/core/engine_adapter/trino.py
index 62808c1b..8834464e 100644
--- a/sqlmesh/core/engine_adapter/trino.py
+++ b/sqlmesh/core/engine_adapter/trino.py
@@ -24,12 +24,13 @@ from sqlmesh.core.engine_adapter.shared import (
     SourceQuery,
     set_catalog,
 )
+from sqlmesh.utils.date import TimeLike
 
 if t.TYPE_CHECKING:
     from trino.dbapi import Connection as TrinoConnection
 
     from sqlmesh.core._typing import SchemaName, TableName
-    from sqlmesh.core.engine_adapter._typing import DF
+    from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF
 
 
 @set_catalog()
@@ -183,3 +184,80 @@ class TrinoEngineAdapter(
                 df[column] = pd.to_datetime(df[column]).map(lambda x: x.isoformat(" "))  # type: ignore
 
         return super()._df_to_source_queries(df, columns_to_types, batch_size, target_table)
+
+    def _build_schema_exp(
+        self,
+        table: exp.Table,
+        columns_to_types: t.Dict[str, exp.DataType],
+        column_descriptions: t.Optional[t.Dict[str, str]] = None,
+        expressions: t.Optional[t.List[exp.PrimaryKey]] = None,
+        is_view: bool = False,
+    ) -> exp.Schema:
+        if self.current_catalog_type == "delta_lake":
+            columns_to_types = self._to_delta_ts(columns_to_types)
+
+        return super()._build_schema_exp(
+            table, columns_to_types, column_descriptions, expressions, is_view
+        )
+
+    def _scd_type_2(
+        self,
+        target_table: TableName,
+        source_table: QueryOrDF,
+        unique_key: t.Sequence[exp.Expression],
+        valid_from_name: str,
+        valid_to_name: str,
+        execution_time: TimeLike,
+        invalidate_hard_deletes: bool = True,
+        updated_at_name: t.Optional[str] = None,
+        check_columns: t.Optional[t.Union[exp.Star, t.Sequence[exp.Column]]] = None,
+        updated_at_as_valid_from: bool = False,
+        execution_time_as_valid_from: bool = False,
+        columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
+        table_description: t.Optional[str] = None,
+        column_descriptions: t.Optional[t.Dict[str, str]] = None,
+    ) -> None:
+        if columns_to_types and self.current_catalog_type == "delta_lake":
+            columns_to_types = self._to_delta_ts(columns_to_types)
+
+        return super()._scd_type_2(
+            target_table,
+            source_table,
+            unique_key,
+            valid_from_name,
+            valid_to_name,
+            execution_time,
+            invalidate_hard_deletes,
+            updated_at_name,
+            check_columns,
+            updated_at_as_valid_from,
+            execution_time_as_valid_from,
+            columns_to_types,
+            table_description,
+            column_descriptions,
+        )
+
+    # delta_lake only supports two timestamp data types. This method converts other
+    # timestamp types to those two for use in DDL statements. Trino/delta automatically
+    # converts the data values to the correct type on write, so we only need to handle
+    # the column types in DDL.
+    # - `timestamp(6)` for non-timezone-aware
+    # - `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]
+    ) -> t.Dict[str, exp.DataType]:
+        ts6 = exp.DataType.build("timestamp(6)")
+        ts3_tz = exp.DataType.build("timestamp(3) with time zone")
+
+        delta_columns_to_types = {
+            k: ts6 if 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
+            for k, v in delta_columns_to_types.items()
+        }
+
+        return delta_columns_to_types


Test patch:
diff --git a/pytest.ini b/pytest.ini
index 38a6c5da..915c8004 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -31,6 +31,7 @@ markers =
     spark: test for Spark
     trino: test for Trino (Hive connector)
     trino_iceberg: test for Trino (Iceberg connector)
+    trino_delta: test for Trino (Delta connector)
 addopts = -n 0 --dist=loadgroup
 
 # Set this to True to enable logging during tests
diff --git a/tests/core/engine_adapter/config.yaml b/tests/core/engine_adapter/config.yaml
index d90ac8bd..fd29007e 100644
--- a/tests/core/engine_adapter/config.yaml
+++ b/tests/core/engine_adapter/config.yaml
@@ -30,6 +30,18 @@ gateways:
       register_comments: true
     state_connection:
       type: duckdb
+  inttest_trino_delta:
+    connection:
+      type: trino
+      host: localhost
+      port: 8080
+      user: admin
+      catalog: datalake_delta
+      http_scheme: http
+      retries: 20
+      register_comments: true
+    state_connection:
+      type: duckdb
   inttest_spark:
     connection:
       type: spark
diff --git a/tests/core/engine_adapter/docker-compose.yaml b/tests/core/engine_adapter/docker-compose.yaml
index 1888fc9e..117a7b1e 100644
--- a/tests/core/engine_adapter/docker-compose.yaml
+++ b/tests/core/engine_adapter/docker-compose.yaml
@@ -60,6 +60,17 @@ services:
     volumes:
       - ./trino/initdb.sql:/docker-entrypoint-initdb.d/initdb.sql
 
+  # A second metastore DB is needed because testing all of hive/iceberg/delta
+  # creates too many connections for a single postgres DB.
+  trino_iceberg_delta_metastore_db:
+    image: postgres
+    hostname: trino_iceberg_delta_metastore_db
+    environment:
+      POSTGRES_USER: hive
+      POSTGRES_PASSWORD: hive
+    volumes:
+      - ./trino/initdb.sql:/docker-entrypoint-initdb.d/initdb.sql
+
   trino-datalake-hive-metastore:
     hostname: trino-datalake-hive-metastore
     image: 'starburstdata/hive:3.1.2-e.15'
@@ -91,13 +102,26 @@ services:
     image: 'starburstdata/hive:3.1.2-e.15'
     environment:
       HIVE_METASTORE_DRIVER: org.postgresql.Driver
-      HIVE_METASTORE_JDBC_URL: jdbc:postgresql://trino_metastore_db:5432/datalake_iceberg_metastore
+      HIVE_METASTORE_JDBC_URL: jdbc:postgresql://trino_iceberg_delta_metastore_db:5432/datalake_iceberg_metastore
       HIVE_METASTORE_USER: hive
       HIVE_METASTORE_PASSWORD: hive
       HIVE_METASTORE_WAREHOUSE_DIR: s3://trino/datalake_iceberg
       <<: *hive_metastore_environments
     depends_on:
-      - trino_metastore_db
+      - trino_iceberg_delta_metastore_db
+
+  trino-datalake-delta-hive-metastore:
+    hostname: trino-datalake-delta-hive-metastore
+    image: 'starburstdata/hive:3.1.2-e.15'
+    environment:
+      HIVE_METASTORE_DRIVER: org.postgresql.Driver
+      HIVE_METASTORE_JDBC_URL: jdbc:postgresql://trino_iceberg_delta_metastore_db:5432/datalake_delta_metastore
+      HIVE_METASTORE_USER: hive
+      HIVE_METASTORE_PASSWORD: hive
+      HIVE_METASTORE_WAREHOUSE_DIR: s3://trino/datalake_delta
+      <<: *hive_metastore_environments
+    depends_on:
+      - trino_iceberg_delta_metastore_db
 
   # Spark Stack
   spark:
@@ -154,8 +178,10 @@ services:
       /usr/bin/mc config --quiet host add myminio http://minio:9000 minio minio123;
       /usr/bin/mc mb --quiet myminio/trino/datalake;
       /usr/bin/mc mb --quiet myminio/trino/datalake_iceberg;
+      /usr/bin/mc mb --quiet myminio/trino/datalake_delta;
       /usr/bin/mc mb --quiet myminio/trino/testing;
       /usr/bin/mc mb --quiet myminio/trino/testing_iceberg;
+      /usr/bin/mc mb --quiet myminio/trino/testing_delta;
       /usr/bin/mc mb --quiet myminio/spark/datalake;
       /usr/bin/mc mb --quiet myminio/spark/testing
       "
diff --git a/tests/core/engine_adapter/test_integration.py b/tests/core/engine_adapter/test_integration.py
index 9c79ee4b..3b97f936 100644
--- a/tests/core/engine_adapter/test_integration.py
+++ b/tests/core/engine_adapter/test_integration.py
@@ -449,6 +449,15 @@ def config() -> Config:
                 pytest.mark.xdist_group("engine_integration_trino_iceberg"),
             ],
         ),
+        pytest.param(
+            "trino_delta",
+            marks=[
+                pytest.mark.docker,
+                pytest.mark.engine,
+                pytest.mark.trino_delta,
+                pytest.mark.xdist_group("engine_integration_trino_delta"),
+            ],
+        ),
         pytest.param(
             "spark",
             marks=[
diff --git a/tests/core/engine_adapter/test_trino.py b/tests/core/engine_adapter/test_trino.py
index 925c581f..b76fa22c 100644
--- a/tests/core/engine_adapter/test_trino.py
+++ b/tests/core/engine_adapter/test_trino.py
@@ -21,6 +21,8 @@ def trino_mocked_engine_adapter(
     def mock_catalog_type(catalog_name):
         if "iceberg" in catalog_name:
             return "iceberg"
+        if "delta" in catalog_name:
+            return "delta"
         return "hive"
 
     mocker.patch(
@@ -40,7 +42,12 @@ def test_set_current_catalog(trino_mocked_engine_adapter: TrinoEngineAdapter):
     ]
 
 
-def test_get_catalog_type(trino_mocked_engine_adapter: TrinoEngineAdapter, mocker: MockerFixture):
+@pytest.mark.trino_iceberg
+@pytest.mark.trino_delta
+@pytest.mark.parametrize("storage_type", ["iceberg", "delta"])
+def test_get_catalog_type(
+    trino_mocked_engine_adapter: TrinoEngineAdapter, mocker: MockerFixture, storage_type: str
+):
     adapter = trino_mocked_engine_adapter
     mocker.patch(
         "sqlmesh.core.engine_adapter.trino.TrinoEngineAdapter.get_current_catalog",
@@ -51,12 +58,13 @@ def test_get_catalog_type(trino_mocked_engine_adapter: TrinoEngineAdapter, mocke
     assert adapter.get_catalog_type("foo") == TrinoEngineAdapter.DEFAULT_CATALOG_TYPE
     assert adapter.get_catalog_type("datalake_hive") == "hive"
     assert adapter.get_catalog_type("datalake_iceberg") == "iceberg"
+    assert adapter.get_catalog_type("datalake_delta") == "delta"
 
     mocker.patch(
         "sqlmesh.core.engine_adapter.trino.TrinoEngineAdapter.get_current_catalog",
-        return_value="system_iceberg",
+        return_value=f"system_{storage_type}",
     )
-    assert adapter.current_catalog_type == "iceberg"
+    assert adapter.current_catalog_type == storage_type
 
 
 def test_get_catalog_type_cached(
@@ -89,16 +97,18 @@ def test_get_catalog_type_cached(
     assert fetchone_mock.call_count == 2
 
 
-def test_partitioned_by_hive(
-    trino_mocked_engine_adapter: TrinoEngineAdapter, mocker: MockerFixture
+@pytest.mark.trino_delta
+@pytest.mark.parametrize("storage_type", ["hive", "delta"])
+def test_partitioned_by_hive_delta(
+    trino_mocked_engine_adapter: TrinoEngineAdapter, mocker: MockerFixture, storage_type: str
 ):
     adapter = trino_mocked_engine_adapter
 
     mocker.patch(
         "sqlmesh.core.engine_adapter.trino.TrinoEngineAdapter.get_current_catalog",
-        return_value="datalake_hive",
+        return_value=f"datalake_{storage_type}",
     )
-    assert adapter.get_catalog_type("datalake_hive") == "hive"
+    assert adapter.get_catalog_type(f"datalake_{storage_type}") == storage_type
 
     columns_to_types = {
         "cola": exp.DataType.build("INT"),
@@ -115,6 +125,7 @@ def test_partitioned_by_hive(
     ]
 
 
+@pytest.mark.trino_iceberg
 def test_partitioned_by_iceberg(
     trino_mocked_engine_adapter: TrinoEngineAdapter, mocker: MockerFixture
 ):
@@ -141,6 +152,7 @@ def test_partitioned_by_iceberg(
     ]
 
 
+@pytest.mark.trino_iceberg
 def test_partitioned_by_iceberg_transforms(
     trino_mocked_engine_adapter: TrinoEngineAdapter, mocker: MockerFixture
 ):
@@ -184,6 +196,7 @@ def test_partitioned_by_iceberg_transforms(
     ]
 
 
+@pytest.mark.trino_iceberg
 def test_partitioned_by_with_multiple_catalogs_same_server(
     trino_mocked_engine_adapter: TrinoEngineAdapter, mocker: MockerFixture
 ):
@@ -194,6 +207,11 @@ def test_partitioned_by_with_multiple_catalogs_same_server(
         "colb": exp.DataType.build("TEXT"),
     }
 
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.trino.TrinoEngineAdapter.get_current_catalog",
+        return_value="system",
+    )
+
     adapter.create_table(
         "datalake.test_schema.test_table", columns_to_types, partitioned_by=[exp.to_column("colb")]
     )
@@ -204,6 +222,11 @@ def test_partitioned_by_with_multiple_catalogs_same_server(
         partitioned_by=[exp.to_column("colb")],
     )
 
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.trino.TrinoEngineAdapter.get_current_catalog",
+        return_value="iceberg",
+    )
+
     adapter.create_table(
         "datalake_iceberg.test_schema.test_table",
         columns_to_types,
@@ -287,17 +310,22 @@ def test_comments_hive(mocker: MockerFixture, make_mocked_engine_adapter: t.Call
     ]
 
 
-def test_comments_iceberg(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable):
+@pytest.mark.trino_iceberg
+@pytest.mark.trino_delta
+@pytest.mark.parametrize("storage_type", ["iceberg", "delta"])
+def test_comments_iceberg_delta(
+    mocker: MockerFixture, make_mocked_engine_adapter: t.Callable, storage_type: str
+):
     adapter = make_mocked_engine_adapter(TrinoEngineAdapter)
 
     current_catalog_mock = mocker.patch(
         "sqlmesh.core.engine_adapter.trino.TrinoEngineAdapter.get_current_catalog"
     )
-    current_catalog_mock.return_value = "iceberg"
+    current_catalog_mock.return_value = storage_type
     catalog_type_mock = mocker.patch(
         "sqlmesh.core.engine_adapter.trino.TrinoEngineAdapter.get_catalog_type"
     )
-    catalog_type_mock.return_value = "iceberg"
+    catalog_type_mock.return_value = storage_type
 
     allowed_table_comment_length = TrinoEngineAdapter.MAX_TABLE_COMMENT_LENGTH
     truncated_table_comment = "a" * allowed_table_comment_length
@@ -348,3 +376,27 @@ def test_comments_iceberg(mocker: MockerFixture, make_mocked_engine_adapter: t.C
         f"""COMMENT ON TABLE "test_table" IS '{long_table_comment}'""",
         f"""COMMENT ON COLUMN "test_table"."a" IS '{long_column_comment}'""",
     ]
+
+
+@pytest.mark.trino_delta
+def test_delta_timestamps(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(TrinoEngineAdapter)
+
+    ts6 = exp.DataType.build("timestamp(6)")
+    ts3_tz = exp.DataType.build("timestamp(3) with time zone")
+
+    columns_to_types = {
+        "ts": exp.DataType.build("TIMESTAMP"),
+        "ts_1": exp.DataType.build("TIMESTAMP(1)"),
+        "ts_tz": exp.DataType.build("TIMESTAMP WITH TIME ZONE"),
+        "ts_tz_1": exp.DataType.build("TIMESTAMP(1) WITH TIME ZONE"),
+    }
+
+    delta_columns_to_types = adapter._to_delta_ts(columns_to_types)
+
+    assert delta_columns_to_types == {
+        "ts": ts6,
+        "ts_1": ts6,
+        "ts_tz": ts3_tz,
+        "ts_tz_1": ts3_tz,
+    }
diff --git a/tests/core/engine_adapter/trino/catalog/datalake_delta.properties b/tests/core/engine_adapter/trino/catalog/datalake_delta.properties
new file mode 100644
index 00000000..a2ee5c99
--- /dev/null
+++ b/tests/core/engine_adapter/trino/catalog/datalake_delta.properties
@@ -0,0 +1,11 @@
+connector.name=delta_lake
+hive.metastore.uri=thrift://trino-datalake-delta-hive-metastore:9083
+hive.metastore-cache-ttl=0s
+hive.metastore-refresh-interval=5s
+hive.metastore-timeout=10s
+hive.s3.endpoint=http://minio:9000
+hive.s3.aws-access-key=minio
+hive.s3.aws-secret-key=minio123
+hive.s3.path-style-access=true
+delta.enable-non-concurrent-writes=true
+delta.hive-catalog-name=datalake_delta
\ No newline at end of file
diff --git a/tests/core/engine_adapter/trino/initdb.sql b/tests/core/engine_adapter/trino/initdb.sql
index 023f3602..ae71ef66 100644
--- a/tests/core/engine_adapter/trino/initdb.sql
+++ b/tests/core/engine_adapter/trino/initdb.sql
@@ -1,4 +1,6 @@
 create database datalake_metastore;
 create database datalake_iceberg_metastore;
+create database datalake_delta_metastore;
 create database testing_metastore;
-create database testing_iceberg_metastore;
\ No newline at end of file
+create database testing_iceberg_metastore;
+create database testing_delta_metastore;
\ No newline at end of file