Source code for airflow.providers.google.cloud.utils.external_token_supplier
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import abc
import time
from functools import wraps
from typing import TYPE_CHECKING, Any
import requests
from google.auth.exceptions import RefreshError
from google.auth.identity_pool import SubjectTokenSupplier
if TYPE_CHECKING:
from google.auth.external_account import SupplierContext
from google.auth.transport import Request
from airflow.utils.log.logging_mixin import LoggingMixin
[docs]def cache_token_decorator(get_subject_token_method):
"""
Cache calls to ``SubjectTokenSupplier`` instances' ``get_token_supplier`` methods.
Different instances of a same SubjectTokenSupplier class with the same attributes
share the OIDC token cache.
:param get_subject_token_method: A method that returns both a token and an integer specifying
the time in seconds until the token expires
See also:
https://googleapis.dev/python/google-auth/latest/reference/google.auth.identity_pool.html#google.auth.identity_pool.SubjectTokenSupplier.get_subject_token
"""
cache = {}
@wraps(get_subject_token_method)
def wrapper(supplier_instance: CacheTokenSupplier, *args, **kwargs) -> str:
"""
Obeys the interface set by ``SubjectTokenSupplier`` for ``get_subject_token`` methods.
:param supplier_instance: the SubjectTokenSupplier instance whose get_subject_token method is being decorated
:return: The token string
"""
nonlocal cache
cache_key = supplier_instance.get_subject_key()
token: dict[str, str | float] = {}
if cache_key not in cache or cache[cache_key]["expiration_time"] < time.monotonic():
supplier_instance.log.info("OIDC token missing or expired")
try:
access_token, expires_in = get_subject_token_method(supplier_instance, *args, **kwargs)
if not isinstance(expires_in, int) or not isinstance(access_token, str):
raise RefreshError # assume error if strange values are provided
except RefreshError:
supplier_instance.log.error("Failed retrieving new OIDC Token from IdP")
raise
expiration_time = time.monotonic() + float(expires_in)
token["access_token"] = access_token
token["expiration_time"] = expiration_time
cache[cache_key] = token
supplier_instance.log.info("New OIDC token retrieved, expires in %s seconds.", expires_in)
return cache[cache_key]["access_token"]
return wrapper
[docs]class CacheTokenSupplier(LoggingMixin, SubjectTokenSupplier):
"""
A superclass for all Subject Token Supplier classes that wish to implement a caching mechanism.
Child classes must implement the ``get_subject_key`` method to generate a string that serves as the cache key,
ensuring that tokens are shared appropriately among instances.
Methods:
get_subject_key: Abstract method to be implemented by child classes. It should return a string that serves as the cache key.
"""
def __init__(self):
super().__init__()
@abc.abstractmethod
[docs] def get_subject_key(self) -> str:
raise NotImplementedError("")
[docs]class ClientCredentialsGrantFlowTokenSupplier(CacheTokenSupplier):
"""
Class that retrieves an OIDC token from an external IdP using OAuth2.0 Client Credentials Grant flow.
This class implements the ``SubjectTokenSupplier`` interface class used by ``google.auth.identity_pool.Credentials``
:params oidc_issuer_url: URL of the IdP that performs OAuth2.0 Client Credentials Grant flow and returns an OIDC token.
:params client_id: Client ID of the application requesting the token
:params client_secret: Client secret of the application requesting the token
:params extra_params_kwargs: Extra parameters to be passed in the payload of the POST request to the `oidc_issuer_url`
See also:
https://googleapis.dev/python/google-auth/latest/reference/google.auth.identity_pool.html#google.auth.identity_pool.SubjectTokenSupplier
"""
def __init__(
self,
oidc_issuer_url: str,
client_id: str,
client_secret: str,
**extra_params_kwargs: Any,
) -> None:
super().__init__()
self.oidc_issuer_url = oidc_issuer_url
self.client_id = client_id
self.client_secret = client_secret
self.extra_params_kwargs = extra_params_kwargs
@cache_token_decorator
[docs] def get_subject_token(self, context: SupplierContext, request: Request) -> tuple[str, int]:
"""Perform Client Credentials Grant flow with IdP and retrieves an OIDC token and expiration time."""
self.log.info("Requesting new OIDC token from external IdP.")
try:
response = requests.post(
self.oidc_issuer_url,
data={
"grant_type": "client_credentials",
"client_id": self.client_id,
"client_secret": self.client_secret,
**self.extra_params_kwargs,
},
)
response.raise_for_status()
except requests.HTTPError as e:
raise RefreshError(str(e))
except requests.ConnectionError as e:
raise RefreshError(str(e))
try:
response_dict = response.json()
except requests.JSONDecodeError:
raise RefreshError(f"Didn't get a json response from {self.oidc_issuer_url}")
# These fields are required
if {"access_token", "expires_in"} - set(response_dict.keys()):
# TODO more information about the error can be provided in the exception by inspecting the response
raise RefreshError(f"No access token returned from {self.oidc_issuer_url}")
return response_dict["access_token"], response_dict["expires_in"]
[docs] def get_subject_key(self) -> str:
"""
Create a cache key using the OIDC issuer URL, client ID, client secret and additional parameters.
Instances with the same credentials will share tokens.
"""
cache_key = (
self.oidc_issuer_url
+ self.client_id
+ self.client_secret
+ ",".join(sorted(self.extra_params_kwargs))
)
return cache_key