| # 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. |
| |
| import base64 |
| import importlib |
| import logging |
| import threading |
| import time |
| from abc import ABC, abstractmethod |
| from functools import cached_property |
| from typing import Any |
| |
| import requests |
| from requests import HTTPError, PreparedRequest, Session |
| from requests.auth import AuthBase |
| |
| from pyiceberg.catalog.rest.response import TokenResponse, _handle_non_200_response |
| from pyiceberg.exceptions import OAuthError |
| |
| AUTH_MANAGER = "auth.manager" |
| |
| COLON = ":" |
| logger = logging.getLogger(__name__) |
| |
| |
| class AuthManager(ABC): |
| """ |
| Abstract base class for Authentication Managers used to supply authorization headers to HTTP clients (e.g. requests.Session). |
| |
| Subclasses must implement the `auth_header` method to return an Authorization header value. |
| """ |
| |
| @abstractmethod |
| def auth_header(self) -> str | None: |
| """Return the Authorization header value, or None if not applicable.""" |
| |
| |
| class NoopAuthManager(AuthManager): |
| """Auth Manager implementation with no auth.""" |
| |
| def auth_header(self) -> str | None: |
| return None |
| |
| |
| class BasicAuthManager(AuthManager): |
| """AuthManager implementation that supports basic password auth.""" |
| |
| def __init__(self, username: str, password: str): |
| credentials = f"{username}:{password}" |
| self._token = base64.b64encode(credentials.encode()).decode() |
| |
| def auth_header(self) -> str: |
| return f"Basic {self._token}" |
| |
| |
| class LegacyOAuth2AuthManager(AuthManager): |
| """Legacy OAuth2 AuthManager implementation. |
| |
| This class exists for backward compatibility, and will be removed in |
| PyIceberg 1.0.0 in favor of OAuth2AuthManager. |
| """ |
| |
| _session: Session |
| _auth_url: str | None |
| _token: str | None |
| _credential: str | None |
| _optional_oauth_params: dict[str, str] | None |
| |
| def __init__( |
| self, |
| session: Session, |
| auth_url: str | None = None, |
| credential: str | None = None, |
| initial_token: str | None = None, |
| optional_oauth_params: dict[str, str] | None = None, |
| ): |
| self._session = session |
| self._auth_url = auth_url |
| self._token = initial_token |
| self._credential = credential |
| self._optional_oauth_params = optional_oauth_params |
| self._refresh_token() |
| |
| def _fetch_access_token(self, credential: str) -> str: |
| if COLON in credential: |
| client_id, client_secret = credential.split(COLON, maxsplit=1) |
| else: |
| client_id, client_secret = None, credential |
| |
| data = {"grant_type": "client_credentials", "client_id": client_id, "client_secret": client_secret} |
| |
| if self._optional_oauth_params: |
| data.update(self._optional_oauth_params) |
| |
| if self._auth_url is None: |
| raise ValueError("Cannot fetch access token from undefined auth_url") |
| |
| response = self._session.post( |
| url=self._auth_url, data=data, headers={**self._session.headers, "Content-type": "application/x-www-form-urlencoded"} |
| ) |
| try: |
| response.raise_for_status() |
| except HTTPError as exc: |
| _handle_non_200_response(exc, {400: OAuthError, 401: OAuthError}) |
| |
| return TokenResponse.model_validate_json(response.text).access_token |
| |
| def _refresh_token(self) -> None: |
| if self._credential is not None: |
| self._token = self._fetch_access_token(self._credential) |
| |
| def auth_header(self) -> str: |
| return f"Bearer {self._token}" |
| |
| |
| class OAuth2TokenProvider: |
| """Thread-safe OAuth2 token provider with token refresh support.""" |
| |
| client_id: str |
| client_secret: str |
| token_url: str |
| scope: str | None |
| refresh_margin: int |
| expires_in: int | None |
| |
| _token: str | None |
| _expires_at: int |
| _lock: threading.Lock |
| |
| def __init__( |
| self, |
| client_id: str, |
| client_secret: str, |
| token_url: str, |
| scope: str | None = None, |
| refresh_margin: int = 60, |
| expires_in: int | None = None, |
| ): |
| self.client_id = client_id |
| self.client_secret = client_secret |
| self.token_url = token_url |
| self.scope = scope |
| self.refresh_margin = refresh_margin |
| self.expires_in = expires_in |
| |
| self._token = None |
| self._expires_at = 0 |
| self._lock = threading.Lock() |
| |
| @cached_property |
| def _client_secret_header(self) -> str: |
| creds = f"{self.client_id}:{self.client_secret}" |
| creds_bytes = creds.encode("utf-8") |
| b64_creds = base64.b64encode(creds_bytes).decode("utf-8") |
| return f"Basic {b64_creds}" |
| |
| def _refresh_token(self) -> None: |
| data = {"grant_type": "client_credentials"} |
| if self.scope: |
| data["scope"] = self.scope |
| |
| response = requests.post(self.token_url, data=data, headers={"Authorization": self._client_secret_header}) |
| response.raise_for_status() |
| result = response.json() |
| |
| self._token = result["access_token"] |
| expires_in = result.get("expires_in", self.expires_in) |
| if expires_in is None: |
| raise ValueError( |
| "The expiration time of the Token must be provided by the Server in the Access Token Response " |
| "in `expires_in` field, or by the PyIceberg Client." |
| ) |
| self._expires_at = time.monotonic() + expires_in - self.refresh_margin |
| |
| def get_token(self) -> str: |
| with self._lock: |
| if not self._token or time.monotonic() >= self._expires_at: |
| self._refresh_token() |
| if self._token is None: |
| raise ValueError("Authorization token is None after refresh") |
| return self._token |
| |
| |
| class OAuth2AuthManager(AuthManager): |
| """Auth Manager implementation that supports OAuth2 as defined in IETF RFC6749.""" |
| |
| def __init__( |
| self, |
| client_id: str, |
| client_secret: str, |
| token_url: str, |
| scope: str | None = None, |
| refresh_margin: int = 60, |
| expires_in: int | None = None, |
| ): |
| self.token_provider = OAuth2TokenProvider( |
| client_id, |
| client_secret, |
| token_url, |
| scope, |
| refresh_margin, |
| expires_in, |
| ) |
| |
| def auth_header(self) -> str: |
| return f"Bearer {self.token_provider.get_token()}" |
| |
| |
| class GoogleAuthManager(AuthManager): |
| """An auth manager that is responsible for handling Google credentials.""" |
| |
| def __init__(self, credentials_path: str | None = None, scopes: list[str] | None = None): |
| """ |
| Initialize GoogleAuthManager. |
| |
| Args: |
| credentials_path: Optional path to Google credentials JSON file. |
| scopes: Optional list of OAuth2 scopes. |
| """ |
| try: |
| import google.auth |
| import google.auth.transport.requests |
| except ImportError as e: |
| raise ImportError("Google Auth libraries not found. Please install 'google-auth'.") from e |
| |
| if credentials_path: |
| self.credentials, _ = google.auth.load_credentials_from_file(credentials_path, scopes=scopes) |
| else: |
| logger.info("Using Google Default Application Credentials") |
| self.credentials, _ = google.auth.default(scopes=scopes) |
| self._auth_request = google.auth.transport.requests.Request() |
| |
| def auth_header(self) -> str: |
| self.credentials.refresh(self._auth_request) |
| return f"Bearer {self.credentials.token}" |
| |
| |
| class EntraAuthManager(AuthManager): |
| """Auth Manager implementation that supports Microsoft Entra ID (Azure AD) authentication. |
| |
| This manager uses the Azure Identity library's DefaultAzureCredential which automatically |
| tries multiple authentication methods including environment variables, managed identity, |
| and Azure CLI. |
| |
| See https://learn.microsoft.com/en-us/azure/developer/python/sdk/authentication/credential-chains |
| for more details on DefaultAzureCredential. |
| """ |
| |
| DEFAULT_SCOPE = "https://storage.azure.com/.default" |
| |
| def __init__( |
| self, |
| scopes: list[str] | None = None, |
| **credential_kwargs: Any, |
| ): |
| """ |
| Initialize EntraAuthManager. |
| |
| Args: |
| scopes: List of OAuth2 scopes. Defaults to ["https://storage.azure.com/.default"]. |
| **credential_kwargs: Arguments passed to DefaultAzureCredential. |
| Supported authentication methods: |
| - Environment Variables: Set AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET |
| - Managed Identity: Works automatically on Azure; for user-assigned, pass managed_identity_client_id |
| - Azure CLI: Works automatically if logged in via `az login` |
| - Workload Identity: Works automatically in AKS with workload identity configured # codespell:ignore aks |
| """ |
| try: |
| from azure.identity import DefaultAzureCredential |
| except ImportError as e: |
| raise ImportError("Azure Identity library not found. Please install with: pip install pyiceberg[entra-auth]") from e |
| |
| self._scopes = scopes or [self.DEFAULT_SCOPE] |
| self._lock = threading.Lock() |
| self._token: str | None = None |
| self._expires_at: float = 0 |
| self._credential = DefaultAzureCredential(**credential_kwargs) |
| |
| def _refresh_token(self) -> None: |
| """Refresh the access token from Azure.""" |
| token = self._credential.get_token(*self._scopes) |
| self._token = token.token |
| # expires_on is a Unix timestamp; add a 60-second margin for safety |
| self._expires_at = token.expires_on - 60 |
| |
| def _get_token(self) -> str: |
| """Get a valid access token, refreshing if necessary.""" |
| with self._lock: |
| if not self._token or time.time() >= self._expires_at: |
| self._refresh_token() |
| if self._token is None: |
| raise ValueError("Failed to obtain Entra access token") |
| return self._token |
| |
| def auth_header(self) -> str: |
| """Return the Authorization header value with a valid Bearer token.""" |
| return f"Bearer {self._get_token()}" |
| |
| |
| class AuthManagerAdapter(AuthBase): |
| """A `requests.auth.AuthBase` adapter for integrating an `AuthManager` into a `requests.Session`. |
| |
| This adapter automatically attaches the appropriate Authorization header to every request. |
| This adapter is useful when working with `requests.Session.auth` |
| and allows reuse of authentication strategies defined by `AuthManager`. |
| This AuthManagerAdapter is only intended to be used against the REST Catalog |
| Server that expects the Authorization Header. |
| """ |
| |
| def __init__(self, auth_manager: AuthManager): |
| """ |
| Initialize AuthManagerAdapter. |
| |
| Args: |
| auth_manager (AuthManager): An instance of an AuthManager subclass. |
| """ |
| self.auth_manager = auth_manager |
| |
| def __call__(self, request: PreparedRequest) -> PreparedRequest: |
| """ |
| Modify the outgoing request to include the Authorization header. |
| |
| Args: |
| request (requests.PreparedRequest): The HTTP request being prepared. |
| |
| Returns: |
| requests.PreparedRequest: The modified request with Authorization header. |
| """ |
| if auth_header := self.auth_manager.auth_header(): |
| request.headers["Authorization"] = auth_header |
| return request |
| |
| |
| class AuthManagerFactory: |
| _registry: dict[str, type["AuthManager"]] = {} |
| |
| @classmethod |
| def register(cls, name: str, auth_manager_class: type["AuthManager"]) -> None: |
| """ |
| Register a string name to a known AuthManager class. |
| |
| Args: |
| name (str): unique name like 'oauth2' to register the AuthManager with |
| auth_manager_class (Type["AuthManager"]): Implementation of AuthManager |
| |
| Returns: |
| None |
| """ |
| cls._registry[name] = auth_manager_class |
| |
| @classmethod |
| def create(cls, class_or_name: str, config: dict[str, Any]) -> AuthManager: |
| """ |
| Create an AuthManager by name or fully-qualified class path. |
| |
| Args: |
| class_or_name (str): Either a name like 'oauth2' or a full class path like 'my.module.CustomAuthManager' |
| config (Dict[str, Any]): Configuration passed to the AuthManager constructor |
| |
| Returns: |
| AuthManager: An instantiated AuthManager subclass |
| """ |
| if class_or_name in cls._registry: |
| manager_cls = cls._registry[class_or_name] |
| else: |
| try: |
| module_path, class_name = class_or_name.rsplit(".", 1) |
| module = importlib.import_module(module_path) |
| manager_cls = getattr(module, class_name) |
| except Exception as err: |
| raise ValueError(f"Could not load AuthManager class for '{class_or_name}'") from err |
| |
| return manager_cls(**config) |
| |
| |
| AuthManagerFactory.register("noop", NoopAuthManager) |
| AuthManagerFactory.register("basic", BasicAuthManager) |
| AuthManagerFactory.register("legacyoauth2", LegacyOAuth2AuthManager) |
| AuthManagerFactory.register("oauth2", OAuth2AuthManager) |
| AuthManagerFactory.register("google", GoogleAuthManager) |
| AuthManagerFactory.register("entra", EntraAuthManager) |
