substitute existing oauth with LegacyOAuth2AuthManager

Author: sungwy

pyiceberg/catalog/rest/__init__.py

@@ -38,7 +38,7 @@
     Catalog,
     PropertiesUpdateSummary,
 )
-from pyiceberg.catalog.rest.auth import AuthManagerAdapter, AuthManagerFactory, LegacyOAuth2AuthManager
+from pyiceberg.catalog.rest.auth import AuthManager, AuthManagerAdapter, AuthManagerFactory, LegacyOAuth2AuthManager
 from pyiceberg.catalog.rest.util import _handle_non_200_response
 from pyiceberg.exceptions import (
     AuthorizationExpiredError,
@@ -244,15 +244,7 @@ def _create_session(self) -> Session:
                 elif ssl_client_cert := ssl_client.get(CERT):
                     session.cert = ssl_client_cert
 
-        auth_config = {
-            "session": self._session,
-            "credential": self.properties.get(CREDENTIAL),
-            "initial_token": self.properties.get(TOKEN),
-            "optional_oauth_params": self._extract_optional_oauth_params(),
-        }
-
-        auth_manager = AuthManagerFactory.create("legacyoauth2", auth_config)
-        session.auth = AuthManagerAdapter(auth_manager)
+        session.auth = AuthManagerAdapter(self._create_legacy_oauth2_auth_manager(session))
         # Set HTTP headers
         self._config_headers(session)
 
@@ -262,6 +254,26 @@ def _create_session(self) -> Session:
 
         return session
 
+    def _create_legacy_oauth2_auth_manager(self, session: Session) -> AuthManager:
+        """Create the LegacyOAuth2AuthManager by fetching required properties.
+
+        This will be deprecated in PyIceberg 1.0
+        """
+        client_credentials = self.properties.get(CREDENTIAL)
+        # We want to call `self.auth_url` only when we are using CREDENTIAL
+        # with the legacy OAUTH2 flow as it will raise a DeprecationWarning
+        auth_url = self.auth_url if client_credentials is not None else None
+
+        auth_config = {
+            "session": session,
+            "auth_url": auth_url,
+            "credential": client_credentials,
+            "initial_token": self.properties.get(TOKEN),
+            "optional_oauth_params": self._extract_optional_oauth_params(),
+        }
+
+        return AuthManagerFactory.create("legacyoauth2", auth_config)
+
     def _check_valid_namespace_identifier(self, identifier: Union[str, Identifier]) -> Identifier:
         """Check if the identifier has at least one element."""
         identifier_tuple = Catalog.identifier_to_tuple(identifier)
@@ -437,8 +449,9 @@ def _refresh_token(self) -> None:
         # Reactive token refresh is atypical - we should proactively refresh tokens in a separate thread
         # instead of retrying on Auth Exceptions. Keeping refresh behavior for the LegacyOAuth2AuthManager
         # for backward compatibility
-        if isinstance(self._session.auth.auth_manager, LegacyOAuth2AuthManager):
-            self._session.auth.auth_manager._refresh_token()
+        auth_manager = self._session.auth.auth_manager  # type: ignore[union-attr]
+        if isinstance(auth_manager, LegacyOAuth2AuthManager):
+            auth_manager._refresh_token()
 
     def _config_headers(self, session: Session) -> None:
         header_properties = get_header_properties(self.properties)

pyiceberg/catalog/rest/auth.py

@@ -56,20 +56,25 @@ def auth_header(self) -> str:
 
 
 class LegacyOAuth2AuthManager(AuthManager):
+    _session: Session
+    _auth_url: Optional[str]
+    _token: Optional[str]
+    _credential: Optional[str]
+    _optional_oauth_params: Optional[Dict[str, str]]
+
     def __init__(
         self,
         session: Session,
-        auth_url: str,
+        auth_url: Optional[str] = None,
         credential: Optional[str] = None,
         initial_token: Optional[str] = None,
         optional_oauth_params: Optional[Dict[str, str]] = None,
     ):
-        self._token: Optional[str] = None
-        self._initial_token = initial_token
-        self._credential = credential
+        self._session = session
         self._auth_url = auth_url
+        self._token = initial_token
+        self._credential = credential
         self._optional_oauth_params = optional_oauth_params
-        self._session = session
         self._refresh_token()
 
     def _fetch_access_token(self, credential: str) -> str:
@@ -83,6 +88,9 @@ def _fetch_access_token(self, credential: str) -> str:
         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"}
         )
@@ -94,9 +102,7 @@ def _fetch_access_token(self, credential: str) -> str:
         return TokenResponse.model_validate_json(response.text).access_token
 
     def _refresh_token(self) -> None:
-        if self._initial_token is not None:
-            self._token = self._initial_token
-        elif self._credential:
+        if self._credential is not None:
             self._token = self._fetch_access_token(self._credential)
 
     def auth_header(self) -> str:
@@ -140,14 +146,21 @@ class AuthManagerFactory:
     _registry: Dict[str, Type["AuthManager"]] = {}
 
     @classmethod
-    def register(cls, name: str, manager_cls: Type["AuthManager"]) -> None:
+    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] = manager_cls
+        cls._registry[name] = auth_manager_class
 
     @classmethod
-    def create(cls, class_or_name: str, config: Dict[str, Any]) -> "AuthManager":
+    def create(cls, class_or_name: str, config: Dict[str, Any]) -> AuthManager:
         """
         Create an AuthManager by name or fully-qualified class path.
 

pyiceberg/catalog/rest/util.py

@@ -14,31 +14,24 @@
 #  KIND, either express or implied.  See the License for the
 #  specific language governing permissions and limitations
 #  under the License.
-from requests import HTTPError
 from json import JSONDecodeError
-from pyiceberg.typedef import IcebergBaseModel
-from pydantic import Field
-from typing import Dict, Type, Optional, Literal
+from typing import Dict, Literal, Optional, Type
+
+from pydantic import Field, ValidationError
+from requests import HTTPError
+
 from pyiceberg.exceptions import (
     AuthorizationExpiredError,
     BadRequestError,
-    CommitFailedException,
-    CommitStateUnknownException,
     ForbiddenError,
-    NamespaceAlreadyExistsError,
-    NamespaceNotEmptyError,
-    NoSuchIdentifierError,
-    NoSuchNamespaceError,
-    NoSuchTableError,
-    NoSuchViewError,
     OAuthError,
     RESTError,
     ServerError,
     ServiceUnavailableError,
-    TableAlreadyExistsError,
     UnauthorizedError,
-    ValidationError
 )
+from pyiceberg.typedef import IcebergBaseModel
+
 
 class TokenResponse(IcebergBaseModel):
     access_token: str = Field()
@@ -48,6 +41,7 @@ class TokenResponse(IcebergBaseModel):
     refresh_token: Optional[str] = Field(default=None)
     scope: Optional[str] = Field(default=None)
 
+
 class ErrorResponseMessage(IcebergBaseModel):
     message: str = Field()
     type: str = Field()
@@ -112,8 +106,6 @@ def _handle_non_200_response(exc: HTTPError, error_handler: Dict[int, Type[Excep
     except ValidationError as e:
         # In the case we don't have a proper response
         errs = ", ".join(err["msg"] for err in e.errors())
-        response = (
-            f"RESTError {exc.response.status_code}: Received unexpected JSON Payload: {exc.response.text}, errors: {errs}"
-        )
+        response = f"RESTError {exc.response.status_code}: Received unexpected JSON Payload: {exc.response.text}, errors: {errs}"
 
     raise exception(response) from exc