add google auth and test

Author: rambleraptor

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
@@ -119,6 +120,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.
 
@@ -197,3 +230,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

@@ -21,9 +21,10 @@
 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,17 @@ def test_basic_auth_header(rest_mock: Mocker) -> None:
     assert len(history) == 1
     actual_headers = history[0].headers
     assert actual_headers["Authorization"] == expected_header
+
+def test_google_auth_header(google_mock: Mocker) -> None:
+    expected_token = "aaaabbb"
+    expected_header = f"Bearer {expected_token}"
+
+    auth_manager = GoogleAuthManager()
+    session = requests.Session()
+    session.auth = AuthManagerAdapter(auth_manager)
+
+    session.get(TEST_URI)
+    history = google_mock.request_history
+    assert len(history) == 2
+    actual_headers = history[1].headers
+    assert actual_headers["Authorization"] == expected_header