Adding support for Google AuthManager

pyiceberg/catalog/rest/auth.py

@@ -18,7 +18,8 @@
 import base64
 import importlib
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional, Type
+import logging
+from typing import Any, Dict, List, Optional, Type
 
 from requests import HTTPError, PreparedRequest, Session
 from requests.auth import AuthBase
@@ -109,6 +110,38 @@ def auth_header(self) -> str:
         return f"Bearer {self._token}"
 
 
+class GoogleAuthManager(AuthManager):
+    """
+    An auth manager that is responsible for handling Google credentials.
+    """
+
+    def __init__(self, credentials_path: Optional[str] = None, scopes: Optional[List[str]] = 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:
+            logging.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) -> Optional[str]:
+        self.credentials.refresh(self._auth_request)
+        return f"Bearer {self.credentials.token}"
+
 class AuthManagerAdapter(AuthBase):
     """A `requests.auth.AuthBase` adapter that integrates an `AuthManager` into a `requests.Session` to automatically attach the appropriate Authorization header to every request.
 
@@ -187,3 +220,4 @@ def create(cls, class_or_name: str, config: Dict[str, Any]) -> AuthManager:
 AuthManagerFactory.register("noop", NoopAuthManager)
 AuthManagerFactory.register("basic", BasicAuthManager)
 AuthManagerFactory.register("legacyoauth2", LegacyOAuth2AuthManager)
+AuthManagerFactory.register("google", GoogleAuthManager)

tests/catalog/test_rest_auth.py

@@ -16,14 +16,15 @@
 #  under the License.
 
 import base64
+from unittest.mock import MagicMock, patch
 
 import pytest
 import requests
 from requests_mock import Mocker
-
-from pyiceberg.catalog.rest.auth import AuthManagerAdapter, BasicAuthManager, NoopAuthManager
+from pyiceberg.catalog.rest.auth import AuthManagerAdapter, BasicAuthManager, GoogleAuthManager, NoopAuthManager
 
 TEST_URI = "https://iceberg-test-catalog/"
+GOOGLE_CREDS_URI = "https://oauth2.googleapis.com/token"
 
 
 @pytest.fixture
@@ -35,6 +36,17 @@ def rest_mock(requests_mock: Mocker) -> Mocker:
     )
     return requests_mock
 
+@pytest.fixture
+def google_mock(requests_mock: Mocker) -> Mocker:
+    requests_mock.post(GOOGLE_CREDS_URI,
+                       json={"access_token": "aaaabbb"},
+                       status_code=200)
+    requests_mock.get(
+        TEST_URI,
+        json={},
+        status_code=200,
+    )
+    return requests_mock
 
 def test_noop_auth_header(rest_mock: Mocker) -> None:
     auth_manager = NoopAuthManager()
@@ -63,3 +75,74 @@ def test_basic_auth_header(rest_mock: Mocker) -> None:
     assert len(history) == 1
     actual_headers = history[0].headers
     assert actual_headers["Authorization"] == expected_header
+
+
+@patch('google.auth.transport.requests.Request')
+@patch('google.auth.default')
+def test_google_auth_manager_default_credentials(mock_google_auth_default: MagicMock, mock_google_request: MagicMock, rest_mock: Mocker) -> None:
+    """Test GoogleAuthManager with default application credentials."""
+    mock_credentials = MagicMock()
+    mock_credentials.token = "test_token"
+    mock_google_auth_default.return_value = (mock_credentials, "test_project")
+
+    auth_manager = GoogleAuthManager()
+    session = requests.Session()
+    session.auth = AuthManagerAdapter(auth_manager)
+    session.get(TEST_URI)
+
+    mock_google_auth_default.assert_called_once_with(scopes=None)
+    mock_credentials.refresh.assert_called_once_with(mock_google_request.return_value)
+    history = rest_mock.request_history
+    assert len(history) == 1
+    actual_headers = history[0].headers
+    assert actual_headers["Authorization"] == "Bearer test_token"
+
+
+@patch('google.auth.transport.requests.Request')
+@patch('google.auth.load_credentials_from_file')
+def test_google_auth_manager_with_credentials_file(mock_load_creds: MagicMock, mock_google_request: MagicMock, rest_mock: Mocker) -> None:
+    """Test GoogleAuthManager with a credentials file path."""
+    mock_credentials = MagicMock()
+    mock_credentials.token = "file_token"
+    mock_load_creds.return_value = (mock_credentials, "test_project_file")
+
+    auth_manager = GoogleAuthManager(credentials_path="/fake/path.json")
+    session = requests.Session()
+    session.auth = AuthManagerAdapter(auth_manager)
+    session.get(TEST_URI)
+
+    mock_load_creds.assert_called_once_with("/fake/path.json", scopes=None)
+    mock_credentials.refresh.assert_called_once_with(mock_google_request.return_value)
+    history = rest_mock.request_history
+    assert len(history) == 1
+    actual_headers = history[0].headers
+    assert actual_headers["Authorization"] == "Bearer file_token"
+
+
+@patch('google.auth.transport.requests.Request')
+@patch('google.auth.load_credentials_from_file')
+def test_google_auth_manager_with_credentials_file_and_scopes(mock_load_creds: MagicMock, mock_google_request: MagicMock, rest_mock: Mocker) -> None:
+    """Test GoogleAuthManager with a credentials file path and scopes."""
+    mock_credentials = MagicMock()
+    mock_credentials.token = "scoped_token"
+    mock_load_creds.return_value = (mock_credentials, "test_project_scoped")
+    scopes = ["https://www.googleapis.com/auth/bigquery"]
+
+    auth_manager = GoogleAuthManager(credentials_path="/fake/path.json", scopes=scopes)
+    session = requests.Session()
+    session.auth = AuthManagerAdapter(auth_manager)
+    session.get(TEST_URI)
+
+    mock_load_creds.assert_called_once_with("/fake/path.json", scopes=scopes)
+    mock_credentials.refresh.assert_called_once_with(mock_google_request.return_value)
+    history = rest_mock.request_history
+    assert len(history) == 1
+    actual_headers = history[0].headers
+    assert actual_headers["Authorization"] == "Bearer scoped_token"
+
+
+def test_google_auth_manager_import_error() -> None:
+    """Test GoogleAuthManager raises ImportError if google-auth is not installed."""
+    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()