From 0417ff073a4b964321cb0c746802a8d41632ef15 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 31 Oct 2025 01:10:14 +0000 Subject: [PATCH 1/2] feat: add OAuth token support for Trino connections Add support for using OAuth Bearer tokens when connecting to Trino databases. When TRINO_OAUTH_TOKEN environment variable is set, the connection will use Bearer authentication instead of basic auth. This enables federated authentication flows where OAuth tokens are managed externally and passed to the SQL execution layer via environment variables. Co-Authored-By: jakub@deepnote.com --- deepnote_toolkit/sql/sql_execution.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/deepnote_toolkit/sql/sql_execution.py b/deepnote_toolkit/sql/sql_execution.py index 7f51e3e..93acc1e 100644 --- a/deepnote_toolkit/sql/sql_execution.py +++ b/deepnote_toolkit/sql/sql_execution.py @@ -107,6 +107,12 @@ class ExecuteSqlError(Exception): params = sql_alchemy_dict.get("params") sql_alchemy_dict["params"] = _build_params_for_bigquery_oauth(params) + requires_trino_oauth = sql_alchemy_dict["url"].startswith("trino://") and dnenv.get_env("TRINO_OAUTH_TOKEN") + + if requires_trino_oauth: + params = sql_alchemy_dict.get("params", {}) + sql_alchemy_dict["params"] = _build_params_for_trino_oauth(params) + # When using key-pair authentication with Snowflake, the private key will be # passed as a base64 encoded string as 'snowflake_private_key'. # @@ -457,6 +463,23 @@ class BigQueryCredentialsError(Exception): return {"connect_args": {"client": client}} +def _build_params_for_trino_oauth(params): + import trino.auth + + oauth_token = dnenv.get_env("TRINO_OAUTH_TOKEN") + if not oauth_token: + raise Exception("TRINO_OAUTH_TOKEN environment variable is not set") + + auth = trino.auth.BearerAuthentication(oauth_token) + + result_params = params.copy() + connect_args = result_params.get("connect_args", {}) + connect_args["auth"] = auth + result_params["connect_args"] = connect_args + + return result_params + + def _sanitize_dataframe_for_parquet(dataframe): """Sanitizes the dataframe so that we can safely call .to_parquet on it""" From 19f0af638d502cb277c9915ff65fb237150a105f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 31 Oct 2025 01:22:33 +0000 Subject: [PATCH 2/2] fix: use JWTAuthentication instead of BearerAuthentication The trino library provides JWTAuthentication for Bearer token auth, not BearerAuthentication. This fixes a runtime error that would occur when trying to use OAuth tokens with Trino connections. Co-Authored-By: jakub@deepnote.com --- deepnote_toolkit/sql/sql_execution.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepnote_toolkit/sql/sql_execution.py b/deepnote_toolkit/sql/sql_execution.py index 93acc1e..5b70790 100644 --- a/deepnote_toolkit/sql/sql_execution.py +++ b/deepnote_toolkit/sql/sql_execution.py @@ -470,7 +470,7 @@ def _build_params_for_trino_oauth(params): if not oauth_token: raise Exception("TRINO_OAUTH_TOKEN environment variable is not set") - auth = trino.auth.BearerAuthentication(oauth_token) + auth = trino.auth.JWTAuthentication(oauth_token) result_params = params.copy() connect_args = result_params.get("connect_args", {})