From eee92e48fb82783a951156706caa19058b2c48e5 Mon Sep 17 00:00:00 2001 From: Sameer Mesiah Date: Mon, 15 Jun 2026 01:45:10 +0100 Subject: [PATCH 1/2] Introduce _SnowflakeOAuthManager to own OAuth token lifecycle management. Move grant type validation, token requests, token caching and token expiration tracking out of SnowflakeHook while preserving the existing public API. SnowflakeHook continues to expose get_oauth_token(), but now delegates OAuth token management to the dedicated helper. --- .../providers/snowflake/hooks/snowflake.py | 232 +++++++++--------- .../unit/snowflake/hooks/test_snowflake.py | 4 +- 2 files changed, 120 insertions(+), 116 deletions(-) diff --git a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py index e25b860524164..56bd9afd2f3f8 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py @@ -86,6 +86,119 @@ def _is_retryable_oauth_error(exception: BaseException) -> bool: return False +class _SnowflakeOAuthManager: + """Encapsulates OAuth token lifecycle management for Snowflake authentication.""" + + def __init__(self): + self._oauth_token: str | None = None + self._oauth_token_expires_at: datetime | None = None + + def validate_grant_type(self, grant_type: str | None) -> str: + """Validate OAuth grant_type.""" + if not grant_type: + raise ValueError("Grant type must be provided for OAuth authentication.") + + if grant_type not in SUPPORTED_GRANT_TYPES: + supported = ", ".join(sorted(SUPPORTED_GRANT_TYPES)) + + raise ValueError(f"Unsupported grant_type '{grant_type}'. Supported values: {supported}") + + return grant_type + + @tenacity.retry( + stop=tenacity.stop_after_attempt(3), + wait=tenacity.wait_exponential(multiplier=1, min=0, max=10), + retry=tenacity.retry_if_exception(_is_retryable_oauth_error), + reraise=True, + ) + def request_oauth_token( + self, + *, + url: str, + data: dict[str, Any], + client_id: str, + client_secret: str, + ): + """ + Execute an OAuth token request. + + Retries automatically on transient failures (ConnectionError, Timeout, 5xx) + via the @tenacity.retry decorator above. + """ + response = requests.post( + url, + data=data, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + auth=HTTPBasicAuth(client_id, client_secret), + timeout=OAUTH_REQUEST_TIMEOUT, + ) + + # Raise HTTPError for non-success responses so retry logic can decide + # whether the failure is retryable. + response.raise_for_status() + return response + + def get_valid_oauth_token( + self, + *, + conn_config: dict[str, Any], + token_endpoint: str | None, + grant_type: str, + ) -> str: + """ + Return a valid OAuth access token. + + This also updates the internal OAuth token cache and token expiry timestamp. + """ + # Check validity using current timestamp. + now = timezone.utcnow() + + if ( + self._oauth_token is not None + and self._oauth_token_expires_at is not None + and now < self._oauth_token_expires_at + ): + return self._oauth_token + + url = token_endpoint or f"https://{conn_config['account']}.snowflakecomputing.com/oauth/token-request" + + data = { + "grant_type": grant_type, + "redirect_uri": conn_config.get("redirect_uri", "https://localhost.com"), + } + + scope = conn_config.get("scope") + + if scope: + data["scope"] = scope + + grant_type = self.validate_grant_type(grant_type) + + if grant_type == "refresh_token": + data |= { + "refresh_token": conn_config["refresh_token"], + } + + response = self.request_oauth_token( + url=url, + data=data, + client_id=conn_config["client_id"], + client_secret=conn_config["client_secret"], + ) + + token = response.json()["access_token"] + expires_in = int(response.json()["expires_in"]) + + # Capture issue timestamp after access token is retrieved. + issued_at = timezone.utcnow() + + # Persist retrieved access token and expiration timestamp. + self._oauth_token = token + self._oauth_token_expires_at = issued_at + timedelta(seconds=max(expires_in - OAUTH_EXPIRY_BUFFER, 0)) + + return token + + class SnowflakeHook(DbApiHook): """ A client to interact with Snowflake. @@ -225,10 +338,7 @@ def __init__(self, *args, **kwargs) -> None: self.client_store_temporary_credential = kwargs.pop("client_store_temporary_credential", None) self.query_ids: list[str] = [] - # Access token and expiration timestamp persisted - # to handle premature expiry. - self._oauth_token: str | None = None - self._oauth_token_expires_at: datetime | None = None + self._oauth = _SnowflakeOAuthManager() def _get_field(self, extra_dict, field_name): backcompat_prefix = "extra__snowflake__" @@ -252,18 +362,6 @@ def _get_field(self, extra_dict, field_name): return extra_dict[field_name] or None return extra_dict.get(backcompat_key) or None - def _validate_grant_type(self, grant_type: str | None) -> str: - """Validate OAuth grant_type.""" - if not grant_type: - raise ValueError("Grant type must be provided for OAuth authentication.") - - if grant_type not in SUPPORTED_GRANT_TYPES: - supported = ", ".join(sorted(SUPPORTED_GRANT_TYPES)) - - raise ValueError(f"Unsupported grant_type '{grant_type}'. Supported values: {supported}") - - return grant_type - @property def account_identifier(self) -> str: """Get snowflake account identifier.""" @@ -292,7 +390,7 @@ def get_oauth_token( if token_endpoint is None: token_endpoint = conn_config.get("token_endpoint") - return self._get_valid_oauth_token( + return self._oauth.get_valid_oauth_token( conn_config=conn_config, token_endpoint=token_endpoint, grant_type=grant_type ) @@ -342,9 +440,9 @@ def _get_conn_params(self) -> dict[str, str | None]: if azure_conn_id: conn_config["token"] = self.get_azure_oauth_token(azure_conn_id) else: - grant_type = self._validate_grant_type(conn_config.get("grant_type")) + grant_type = self._oauth.validate_grant_type(conn_config.get("grant_type")) - conn_config["token"] = self._get_valid_oauth_token( + conn_config["token"] = self._oauth.get_valid_oauth_token( conn_config=conn_config, token_endpoint=conn_config.get("token_endpoint"), grant_type=grant_type, @@ -361,8 +459,7 @@ def _get_static_conn_params(self) -> dict[str, str | None]: Return static Snowflake connection parameters. These parameters are cached for the lifetime of the hook and exclude - time-sensitive values such as OAuth access tokens. This is used in - ``_get_valid_oauth_token()`` and ``get_conn_params()``. + time-sensitive values such as OAuth access tokens. """ conn = self.get_connection(self.get_conn_id()) extra_dict = conn.extra_dejson @@ -491,99 +588,6 @@ def _get_static_conn_params(self) -> dict[str, str | None]: return conn_config - def _get_valid_oauth_token( - self, - *, - conn_config: dict[str, Any], - token_endpoint: str | None, - grant_type: str, - ) -> str: - """ - Return a valid OAuth access token. - - This also updates the internal OAuth token cache and token expiry timestamp. - """ - # Check validity using current timestamp. - now = timezone.utcnow() - - if ( - self._oauth_token is not None - and self._oauth_token_expires_at is not None - and now < self._oauth_token_expires_at - ): - return self._oauth_token - - url = token_endpoint or f"https://{conn_config['account']}.snowflakecomputing.com/oauth/token-request" - - data = { - "grant_type": grant_type, - "redirect_uri": conn_config.get("redirect_uri", "https://localhost.com"), - } - - scope = conn_config.get("scope") - - if scope: - data["scope"] = scope - - grant_type = self._validate_grant_type(grant_type) - - if grant_type == "refresh_token": - data |= { - "refresh_token": conn_config["refresh_token"], - } - - response = self._request_oauth_token( - url=url, - data=data, - client_id=conn_config["client_id"], - client_secret=conn_config["client_secret"], - ) - - token = response.json()["access_token"] - expires_in = int(response.json()["expires_in"]) - - # Capture issue timestamp after access token is retrieved. - issued_at = timezone.utcnow() - - # Persist retrieved access token and expiration timestamp. - self._oauth_token = token - self._oauth_token_expires_at = issued_at + timedelta(seconds=max(expires_in - OAUTH_EXPIRY_BUFFER, 0)) - - return token - - @tenacity.retry( - stop=tenacity.stop_after_attempt(3), - wait=tenacity.wait_exponential(multiplier=1, min=0, max=10), - retry=tenacity.retry_if_exception(_is_retryable_oauth_error), - reraise=True, - ) - def _request_oauth_token( - self, - *, - url: str, - data: dict[str, Any], - client_id: str, - client_secret: str, - ): - """ - Execute a single OAuth token request. - - Performs one HTTP call and raises ``HTTPError`` for 4xx and 5xx responses. - Retry behavior is handled by the caller. - """ - response = requests.post( - url, - data=data, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - auth=HTTPBasicAuth(client_id, client_secret), - timeout=OAUTH_REQUEST_TIMEOUT, - ) - - # Raise HTTPError for non-success responses so retry logic can decide - # whether the failure is retryable. - response.raise_for_status() - return response - def get_private_key(self) -> PrivateKeyTypes | None: """Get the private key from snowflake connection.""" conn = self.get_connection(self.get_conn_id()) diff --git a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py index 5b96484933e07..9ffc3332b0bab 100644 --- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py @@ -1314,9 +1314,9 @@ def test_validate_grant_type(self, grant_type, expected, match): if expected is ValueError: with pytest.raises(ValueError, match=match): - hook._validate_grant_type(grant_type) + hook._oauth.validate_grant_type(grant_type) else: - assert hook._validate_grant_type(grant_type) == expected + assert hook._oauth.validate_grant_type(grant_type) == expected @mock.patch("airflow.providers.snowflake.hooks.snowflake.HTTPBasicAuth") @mock.patch("requests.post") From 4ed50ee6d07308d14d1aa9c14d5088bd613a879d Mon Sep 17 00:00:00 2001 From: Sameer Mesiah Date: Tue, 16 Jun 2026 18:26:49 +0100 Subject: [PATCH 2/2] Adjust comment in the test file. --- .../snowflake/tests/unit/snowflake/hooks/test_snowflake.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py index 9ffc3332b0bab..d0fdb49b5b5cc 100644 --- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py @@ -1599,7 +1599,7 @@ def test_oauth_token_refresh_after_expiry(self, mock_timezone_utcnow, mock_reque t0 = datetime(2025, 1, 1, 12, 0, tzinfo=timezone.utc) - # _get_valid_oauth_token calls utcnow twice per refresh: + # get_valid_oauth_token from _SnowflakeOAuthManager calls utcnow twice per refresh: # 1) validity check # 2) issued_at mock_timezone_utcnow.side_effect = [