phanikumv commented on code in PR #68549:
URL: https://github.com/apache/airflow/pull/68549#discussion_r3411176517
##########
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 a single OAuth token request.
+
+ Performs one HTTP call and raises ``HTTPError`` for 4xx and 5xx
responses.
+ Retry behavior is handled by the caller.
+ """
Review Comment:
nit: Outdated docstring, this is not a one HTTP call anymore, due to
`@tenacity.retry`
```suggestion
"""
Execute an OAuth token request.
Retries automatically on transient failures (ConnectionError,
Timeout, 5xx)
via the @tenacity.retry decorator above.
"""
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
