From c6081a73069c8e52827144bdbb4559f390825107 Mon Sep 17 00:00:00 2001 From: Jesse Hodges Date: Thu, 15 Jan 2026 18:20:35 -0600 Subject: [PATCH 1/2] expose a source option for trino --- docs/integrations/engines/trino.md | 1 + .../models/jesse_sqlmesh_source_test.sql | 15 +++++++++ sqlmesh/core/config/connection.py | 4 ++- tests/core/engine_adapter/test_trino.py | 4 +++ tests/core/test_config.py | 33 +++++++++++++++++++ 5 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 examples/trino-test-proj/models/jesse_sqlmesh_source_test.sql diff --git a/docs/integrations/engines/trino.md b/docs/integrations/engines/trino.md index ec1139e20d..db732f0cc1 100644 --- a/docs/integrations/engines/trino.md +++ b/docs/integrations/engines/trino.md @@ -90,6 +90,7 @@ hive.metastore.glue.default-warehouse-dir=s3://my-bucket/ | `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 | +| `source` | Value to send as Trino's `source` field for query attribution / auditing. Default: `sqlmesh`. | string | 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 | diff --git a/examples/trino-test-proj/models/jesse_sqlmesh_source_test.sql b/examples/trino-test-proj/models/jesse_sqlmesh_source_test.sql new file mode 100644 index 0000000000..bd524e4bd6 --- /dev/null +++ b/examples/trino-test-proj/models/jesse_sqlmesh_source_test.sql @@ -0,0 +1,15 @@ +MODEL ( + name css_data.scratch.jesse_source_test, + owner 'jesse.hodges', + description 'Simple test model to verify potential source addition to sqlmesh', + kind VIEW, + cron '@daily' +); + +-- Simple test query that returns static data +SELECT + CURRENT_TIMESTAMP AS test_timestamp, + 'jesse_source_test' AS model_name, + 42 AS test_number, + 666 AS another_test_number_2, + TRUE AS is_active; diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 638f0c28c8..4e11fc626f 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -1888,6 +1888,7 @@ class TrinoConnectionConfig(ConnectionConfig): client_certificate: t.Optional[str] = None client_private_key: t.Optional[str] = None cert: t.Optional[str] = None + source: str = "sqlmesh" # SQLMesh options schema_location_mapping: t.Optional[dict[re.Pattern, str]] = None @@ -1984,6 +1985,7 @@ def _connection_kwargs_keys(self) -> t.Set[str]: "port", "catalog", "roles", + "source", "http_scheme", "http_headers", "session_properties", @@ -2041,7 +2043,7 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: "user": self.impersonation_user or self.user, "max_attempts": self.retries, "verify": self.cert if self.cert is not None else self.verify, - "source": "sqlmesh", + "source": self.source, } @property diff --git a/tests/core/engine_adapter/test_trino.py b/tests/core/engine_adapter/test_trino.py index a3c67eb023..1bfe82b858 100644 --- a/tests/core/engine_adapter/test_trino.py +++ b/tests/core/engine_adapter/test_trino.py @@ -412,6 +412,8 @@ def test_timestamp_mapping(): catalog="catalog", ) + assert config._connection_factory_with_kwargs.keywords["source"] == "sqlmesh" + adapter = config.create_engine_adapter() assert adapter.timestamp_mapping is None @@ -419,11 +421,13 @@ def test_timestamp_mapping(): user="user", host="host", catalog="catalog", + source="my_source", timestamp_mapping={ "TIMESTAMP": "TIMESTAMP(6)", "TIMESTAMP(3)": "TIMESTAMP WITH TIME ZONE", }, ) + assert config._connection_factory_with_kwargs.keywords["source"] == "my_source" adapter = config.create_engine_adapter() assert adapter.timestamp_mapping is not None assert adapter.timestamp_mapping[exp.DataType.build("TIMESTAMP")] == exp.DataType.build( diff --git a/tests/core/test_config.py b/tests/core/test_config.py index d0fad16e76..f3a0de6672 100644 --- a/tests/core/test_config.py +++ b/tests/core/test_config.py @@ -862,6 +862,39 @@ def test_trino_schema_location_mapping_syntax(tmp_path): assert len(conn.schema_location_mapping) == 2 +def test_trino_source_option(tmp_path): + config_path = tmp_path / "config_trino_source.yaml" + with open(config_path, "w", encoding="utf-8") as fd: + fd.write( + """ + gateways: + trino: + connection: + type: trino + user: trino + host: trino + catalog: trino + source: my_sqlmesh_source + + default_gateway: trino + + model_defaults: + dialect: trino + """ + ) + + config = load_config_from_paths( + Config, + project_paths=[config_path], + ) + + from sqlmesh.core.config.connection import TrinoConnectionConfig + + conn = config.gateways["trino"].connection + assert isinstance(conn, TrinoConnectionConfig) + assert conn.source == "my_sqlmesh_source" + + def test_gcp_postgres_ip_and_scopes(tmp_path): config_path = tmp_path / "config_gcp_postgres.yaml" with open(config_path, "w", encoding="utf-8") as fd: From 2776c962738e5ca0959aa901e8957cf92cc95eb7 Mon Sep 17 00:00:00 2001 From: Jesse Hodges Date: Fri, 16 Jan 2026 09:48:28 -0600 Subject: [PATCH 2/2] rm example model --- .../models/jesse_sqlmesh_source_test.sql | 15 --------------- 1 file changed, 15 deletions(-) delete mode 100644 examples/trino-test-proj/models/jesse_sqlmesh_source_test.sql diff --git a/examples/trino-test-proj/models/jesse_sqlmesh_source_test.sql b/examples/trino-test-proj/models/jesse_sqlmesh_source_test.sql deleted file mode 100644 index bd524e4bd6..0000000000 --- a/examples/trino-test-proj/models/jesse_sqlmesh_source_test.sql +++ /dev/null @@ -1,15 +0,0 @@ -MODEL ( - name css_data.scratch.jesse_source_test, - owner 'jesse.hodges', - description 'Simple test model to verify potential source addition to sqlmesh', - kind VIEW, - cron '@daily' -); - --- Simple test query that returns static data -SELECT - CURRENT_TIMESTAMP AS test_timestamp, - 'jesse_source_test' AS model_name, - 42 AS test_number, - 666 AS another_test_number_2, - TRUE AS is_active;