diff --git a/.changes/unreleased/Features-20260329-152534.yaml b/.changes/unreleased/Features-20260329-152534.yaml new file mode 100644 index 00000000..0c1614a2 --- /dev/null +++ b/.changes/unreleased/Features-20260329-152534.yaml @@ -0,0 +1,7 @@ +kind: Features +body: Add gcp_adc authentication method with token refreshing +time: 2026-03-29T15:25:34.597868-04:00 +custom: + Author: artemboiko-data + Issue: "" + PR: "512" diff --git a/dbt/adapters/trino/connections.py b/dbt/adapters/trino/connections.py index 3b89aa20..4994ed53 100644 --- a/dbt/adapters/trino/connections.py +++ b/dbt/adapters/trino/connections.py @@ -46,6 +46,8 @@ def _create_trino_profile(cls, profile): return TrinoOauthCredentials elif method == "oauth_console": return TrinoOauthConsoleCredentials + elif method == "gcp_adc": + return TrinoGcpAdcCredentials return TrinoNoneCredentials @classmethod @@ -91,6 +93,9 @@ def _connection_keys(self): def trino_auth(self) -> Optional[trino.auth.Authentication]: pass + def http_session(self): + return None + @dataclass class TrinoNoneCredentials(TrinoCredentials): @@ -311,6 +316,49 @@ def trino_auth(self): return self.OAUTH +@dataclass +class TrinoGcpAdcCredentials(TrinoCredentials): + host: str + port: Port + user: Optional[str] = None + client_tags: Optional[List[str]] = None + roles: Optional[Dict[str, str]] = None + cert: Optional[Union[str, bool]] = None + http_headers: Optional[Dict[str, str]] = None + session_properties: Dict[str, Any] = field(default_factory=dict) + prepared_statements_enabled: bool = PREPARED_STATEMENTS_ENABLED_DEFAULT + retries: Optional[int] = trino.constants.DEFAULT_MAX_ATTEMPTS + timezone: Optional[str] = None + suppress_cert_warning: Optional[bool] = None + scopes: List[str] = field( + default_factory=lambda: ["openid", "https://www.googleapis.com/auth/userinfo.email"] + ) + + @property + def http_scheme(self): + return HttpScheme.HTTPS + + @property + def method(self): + return "gcp_adc" + + def trino_auth(self): + return trino.constants.DEFAULT_AUTH + + def http_session(self): + try: + import google.auth + from google.auth.transport.requests import AuthorizedSession + except ImportError: + raise DbtRuntimeError( + "The 'google-auth' library is required for the gcp_adc method. " + "Please install it with: pip install google-auth" + ) + + credentials, project_id = google.auth.default(scopes=self.scopes) + return AuthorizedSession(credentials) + + class ConnectionWrapper(object): """Wrap a Trino connection in a way that accomplishes two tasks: @@ -506,6 +554,7 @@ def open(cls, connection): http_headers=credentials.http_headers, session_properties=credentials.session_properties, auth=credentials.trino_auth(), + http_session=credentials.http_session(), max_attempts=credentials.retries, isolation_level=IsolationLevel.AUTOCOMMIT, source=f"dbt-trino-{version}", diff --git a/setup.py b/setup.py index 4d5ccab3..e551246c 100644 --- a/setup.py +++ b/setup.py @@ -82,6 +82,9 @@ def _dbt_trino_version(): # add dbt-core to ensure backwards compatibility of installation, this is not a functional dependency "dbt-core>=1.8.0", ], + extras_require={ + "gcp": ["google-auth>=2.49.1"], + }, zip_safe=False, classifiers=[ "Development Status :: 5 - Production/Stable", diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py index c389bc84..7b440853 100644 --- a/tests/unit/test_adapter.py +++ b/tests/unit/test_adapter.py @@ -22,6 +22,7 @@ TrinoNoneCredentials, TrinoOauthConsoleCredentials, TrinoOauthCredentials, + TrinoGcpAdcCredentials, ) from .utils import config_from_parts_or_dicts, mock_connection @@ -446,6 +447,37 @@ def test_oauth_console_authentication(self): self.assertEqual(credentials.timezone, "UTC") self.assertEqual(credentials.suppress_cert_warning, False) + def test_gcp_adc_authentication(self): + connection = self.acquire_connection_with_profile( + { + "type": "trino", + "catalog": "trinodb", + "host": "database", + "port": 5439, + "method": "gcp_adc", + "schema": "dbt_test_schema", + "cert": "/path/to/cert", + "client_tags": ["dev", "gcp_adc"], + "http_headers": {"X-Trino-Client-Info": "dbt-trino"}, + "session_properties": { + "query_max_run_time": "4h", + "exchange_compression": True, + }, + "timezone": "UTC", + "suppress_cert_warning": False, + } + ) + credentials = connection.credentials + self.assertIsInstance(credentials, TrinoGcpAdcCredentials) + self.assert_default_connection_credentials(credentials) + self.assertEqual(credentials.http_scheme, HttpScheme.HTTPS) + self.assertEqual(credentials.cert, "/path/to/cert") + self.assertEqual(connection.credentials.prepared_statements_enabled, True) + self.assertEqual(credentials.client_tags, ["dev", "gcp_adc"]) + self.assertEqual(credentials.timezone, "UTC") + self.assertEqual(credentials.suppress_cert_warning, False) + self.assertEqual(credentials.scopes, ["openid", "https://www.googleapis.com/auth/userinfo.email"]) + class TestPreparedStatementsEnabled(TestCase): def setup_profile(self, credentials):