Lee-W commented on code in PR #73932:
URL: https://github.com/apache/airflow/pull/73932#discussion_r4178008064
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py:
##########
@@ -230,54 +231,26 @@ def execute_query(
def get_headers(self) -> dict[str, Any]:
"""Form auth headers based on OAuth token, PAT, or JWT token from
private key."""
- conn_config = self._get_conn_params()
-
- # _get_conn_params() already fetched the OAuth access token for any
grant type or azure_conn_id.
- if conn_config.get("authenticator") == "oauth":
- return {
- "Content-Type": "application/json",
- "Authorization": f"Bearer {conn_config['token']}",
- "Accept": "application/json",
- "User-Agent": "snowflakeSQLAPI/1.0",
- "X-Snowflake-Authorization-Token-Type": "OAUTH",
- }
-
- # Use PAT (Programmatic Access Token) when authenticator is set to
programmatic_access_token
- if conn_config.get("authenticator") == "programmatic_access_token":
- pat = conn_config.get("password")
- if not pat:
- raise ValueError(
- "Programmatic Access Token (PAT) authentication requires
the connection password "
- "field to contain the PAT token value."
- )
- return {
- "Content-Type": "application/json",
- "Authorization": f"Bearer {pat}",
- "Accept": "application/json",
- "User-Agent": "snowflakeSQLAPI/1.0",
- "X-Snowflake-Authorization-Token-Type":
"PROGRAMMATIC_ACCESS_TOKEN",
- }
-
- # Fall back to JWT token from the connection details and the private
key
- if not self.private_key:
- self.private_key = self.get_private_key()
-
- token = JWTGenerator(
- conn_config["account"], # type: ignore[arg-type]
- conn_config["user"], # type: ignore[arg-type]
- private_key=self.private_key,
- lifetime=self.token_life_time,
- renewal_delay=self.token_renewal_delta,
- ).get_token()
-
+ if self._rest_token_provider is None:
Review Comment:
Each partition request now calls `get_headers()` and reuses the cached
provider.
--
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]
