add entra auth manager impl

Author: kevinjqliu

pyiceberg/catalog/rest/auth.py

@@ -249,6 +249,68 @@ def auth_header(self) -> str:
         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`.
 
@@ -330,3 +392,4 @@ def create(cls, class_or_name: str, config: dict[str, Any]) -> AuthManager:
 AuthManagerFactory.register("legacyoauth2", LegacyOAuth2AuthManager)
 AuthManagerFactory.register("oauth2", OAuth2AuthManager)
 AuthManagerFactory.register("google", GoogleAuthManager)
+AuthManagerFactory.register("entra", EntraAuthManager)

tests/catalog/test_rest_auth.py

@@ -22,7 +22,7 @@
 import requests
 from requests_mock import Mocker
 
-from pyiceberg.catalog.rest.auth import AuthManagerAdapter, BasicAuthManager, GoogleAuthManager, NoopAuthManager
+from pyiceberg.catalog.rest.auth import AuthManagerAdapter, BasicAuthManager, EntraAuthManager, GoogleAuthManager, NoopAuthManager
 
 TEST_URI = "https://iceberg-test-catalog/"
 GOOGLE_CREDS_URI = "https://oauth2.googleapis.com/token"
@@ -153,3 +153,80 @@ def test_google_auth_manager_import_error() -> None:
     with patch.dict("sys.modules", {"google.auth": None, "google.auth.transport.requests": None}):
         with pytest.raises(ImportError, match="Google Auth libraries not found. Please install 'google-auth'."):
             GoogleAuthManager()
+
+
+@patch("azure.identity.DefaultAzureCredential")
+def test_entra_auth_manager_default_credential(mock_default_cred: MagicMock, rest_mock: Mocker) -> None:
+    """Test EntraAuthManager with DefaultAzureCredential."""
+    mock_credential_instance = MagicMock()
+    mock_token = MagicMock()
+    mock_token.token = "entra_default_token"
+    mock_token.expires_on = 9999999999  # Far future timestamp
+    mock_credential_instance.get_token.return_value = mock_token
+    mock_default_cred.return_value = mock_credential_instance
+
+    auth_manager = EntraAuthManager()
+    session = requests.Session()
+    session.auth = AuthManagerAdapter(auth_manager)
+    session.get(TEST_URI)
+
+    mock_default_cred.assert_called_once_with()
+    mock_credential_instance.get_token.assert_called_once_with("https://storage.azure.com/.default")
+    history = rest_mock.request_history
+    assert len(history) == 1
+    actual_headers = history[0].headers
+    assert actual_headers["Authorization"] == "Bearer entra_default_token"
+
+
+@patch("azure.identity.DefaultAzureCredential")
+def test_entra_auth_manager_with_managed_identity_client_id(mock_default_cred: MagicMock, rest_mock: Mocker) -> None:
+    """Test EntraAuthManager with managed_identity_client_id passed to DefaultAzureCredential."""
+    mock_credential_instance = MagicMock()
+    mock_token = MagicMock()
+    mock_token.token = "entra_mi_token"
+    mock_token.expires_on = 9999999999
+    mock_credential_instance.get_token.return_value = mock_token
+    mock_default_cred.return_value = mock_credential_instance
+
+    auth_manager = EntraAuthManager(managed_identity_client_id="user-assigned-client-id")
+    session = requests.Session()
+    session.auth = AuthManagerAdapter(auth_manager)
+    session.get(TEST_URI)
+
+    mock_default_cred.assert_called_once_with(managed_identity_client_id="user-assigned-client-id")
+    mock_credential_instance.get_token.assert_called_once_with("https://storage.azure.com/.default")
+    history = rest_mock.request_history
+    assert len(history) == 1
+    actual_headers = history[0].headers
+    assert actual_headers["Authorization"] == "Bearer entra_mi_token"
+
+
+@patch("azure.identity.DefaultAzureCredential")
+def test_entra_auth_manager_custom_scopes(mock_default_cred: MagicMock, rest_mock: Mocker) -> None:
+    """Test EntraAuthManager with custom scopes."""
+    mock_credential_instance = MagicMock()
+    mock_token = MagicMock()
+    mock_token.token = "entra_custom_scope_token"
+    mock_token.expires_on = 9999999999
+    mock_credential_instance.get_token.return_value = mock_token
+    mock_default_cred.return_value = mock_credential_instance
+
+    custom_scopes = ["https://datalake.azure.net/.default", "https://storage.azure.com/.default"]
+    auth_manager = EntraAuthManager(scopes=custom_scopes)
+    session = requests.Session()
+    session.auth = AuthManagerAdapter(auth_manager)
+    session.get(TEST_URI)
+
+    mock_default_cred.assert_called_once_with()
+    mock_credential_instance.get_token.assert_called_once_with(*custom_scopes)
+    history = rest_mock.request_history
+    assert len(history) == 1
+    actual_headers = history[0].headers
+    assert actual_headers["Authorization"] == "Bearer entra_custom_scope_token"
+
+
+def test_entra_auth_manager_import_error() -> None:
+    """Test EntraAuthManager raises ImportError if azure-identity is not installed."""
+    with patch.dict("sys.modules", {"azure.identity": None}):
+        with pytest.raises(ImportError, match="Azure Identity library not found"):
+            EntraAuthManager()