From d47566f3a1e06b4c12cd94e1b9161aea3510457c Mon Sep 17 00:00:00 2001 From: Dan Flippo Date: Fri, 10 Jan 2025 14:31:10 -0500 Subject: [PATCH] Updated to allow the token_file_path parameter Snowflake Container Services (SPCS) provides an OAuth token file, `/snowflake/session/token`, to containers to allow them to log into Snowflake without needing a service account. The token in this file is refreshed automatically by the SPCS service and expires after a few minutes. The Snowflake Python driver added the `token_file_path` parameter to make it easier for clients to connect without having to read the oauth token from the file every time a new connection is needed. These changes have been tested to work in SPCS. --- dbt/adapters/snowflake/connections.py | 49 +++++++++++++++++---------- 1 file changed, 31 insertions(+), 18 deletions(-) diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py index fc2c09c19..b27ab4917 100644 --- a/dbt/adapters/snowflake/connections.py +++ b/dbt/adapters/snowflake/connections.py @@ -96,6 +96,7 @@ class SnowflakeCredentials(Credentials): private_key_path: Optional[str] = None private_key_passphrase: Optional[str] = None token: Optional[str] = None + token_file_path: Optional[str] = None oauth_client_id: Optional[str] = None oauth_client_secret: Optional[str] = None query_tag: Optional[str] = None @@ -114,7 +115,9 @@ class SnowflakeCredentials(Credentials): reuse_connections: Optional[bool] = None def __post_init__(self): - if self.authenticator != "oauth" and (self.oauth_client_secret or self.oauth_client_id): + if self.authenticator != "oauth" and ( + self.oauth_client_secret or self.oauth_client_id or self.token_file_path + ): # the user probably forgot to set 'authenticator' like I keep doing warn_or_error( AdapterEventWarning( @@ -133,8 +136,9 @@ def __post_init__(self): ) ) - if not self.user: + if not (self.user or self.token_file_path): # The user attribute is only optional if 'authenticator' is 'jwt' or 'oauth' + # It is also optional when using an OAuth token file in Snowpark Container Services warn_or_error( AdapterEventError(base_msg="Invalid profile: 'user' is a required property.") ) @@ -179,6 +183,7 @@ def _connection_keys(self): "retry_all", "insecure_mode", "reuse_connections", + "token_file_path", ) def auth_args(self): @@ -202,25 +207,33 @@ def auth_args(self): if self.authenticator: result["authenticator"] = self.authenticator if self.authenticator == "oauth": - token = self.token - # if we have a client ID/client secret, the token is a refresh - # token, not an access token - if self.oauth_client_id and self.oauth_client_secret: - token = self._get_access_token() - elif self.oauth_client_id: - warn_or_error( - AdapterEventWarning( - base_msg="Invalid profile: got an oauth_client_id, but not an oauth_client_secret!" + # If the token_file_path is provided we ignore the token parameter + if self.token_file_path: + if not os.path.isfile(self.token_file_path): + raise DbtInternalError( + f"The token_file_path file does not exist: {self.token_file_path}" ) - ) - elif self.oauth_client_secret: - warn_or_error( - AdapterEventWarning( - base_msg="Invalid profile: got an oauth_client_secret, but not an oauth_client_id!" + result["token_file_path"] = self.token_file_path + else: + token = self.token + # if we have a client ID/client secret, the token is a refresh + # token, not an access token + if self.oauth_client_id and self.oauth_client_secret: + token = self._get_access_token() + elif self.oauth_client_id: + warn_or_error( + AdapterEventWarning( + base_msg="Invalid profile: got an oauth_client_id, but not an oauth_client_secret!" + ) + ) + elif self.oauth_client_secret: + warn_or_error( + AdapterEventWarning( + base_msg="Invalid profile: got an oauth_client_secret, but not an oauth_client_id!" + ) ) - ) - result["token"] = token + result["token"] = token elif self.authenticator == "jwt": # If authenticator is 'jwt', then the 'token' value should be used