refactor(rest): add typed auth config parsing helpers

Author: Gayathri Srividya Rajavarapu

pyiceberg/catalog/rest/__init__.py

@@ -34,7 +34,7 @@
 from pydantic import ConfigDict, Field, TypeAdapter, field_validator
 from requests import HTTPError, Session
 from tenacity import RetryCallState, retry, retry_if_exception_type, stop_after_attempt
-from typing_extensions import override
+from typing_extensions import NotRequired, TypedDict, override
 
 from pyiceberg import __version__
 from pyiceberg.catalog import BOTOCORE_SESSION, TOKEN, URI, WAREHOUSE_LOCATION, Catalog, PropertiesUpdateSummary
@@ -403,6 +403,17 @@ class ListViewsResponse(IcebergBaseModel):
 _PLANNING_RESPONSE_ADAPTER = TypeAdapter(PlanningResponse)
 
 
+class ParsedAuthConfig(TypedDict):
+    auth_type: str
+    auth_manager_name: str
+    auth_type_config: dict[str, Any]
+
+
+class AuthConfigEnvelope(TypedDict):
+    type: str
+    impl: NotRequired[str]
+
+
 def _get_auth_manager_class(class_or_name: str) -> type[AuthManager]:
     if class_or_name in AuthManagerFactory._registry:
         return AuthManagerFactory._registry[class_or_name]
@@ -461,6 +472,65 @@ def _coerce_auth_config_values(class_or_name: str, config: dict[str, Any]) -> di
     return {key: _coerce_auth_option_value(key, value, hints.get(key, Any)) for key, value in config.items()}
 
 
+def _load_auth_config_from_properties(properties: Properties) -> AuthConfigEnvelope | dict[str, Any] | None:
+    raw_auth = properties.get(AUTH)
+    if isinstance(raw_auth, str):
+        try:
+            decoded_auth = json.loads(raw_auth)
+        except json.JSONDecodeError as e:
+            raise ValueError("Failed to parse auth configuration as JSON") from e
+        if decoded_auth is not None and not isinstance(decoded_auth, dict):
+            raise ValueError("auth configuration must be a dictionary")
+        return decoded_auth
+
+    if raw_auth is not None:
+        if not isinstance(raw_auth, dict):
+            raise ValueError("auth configuration must be a dictionary")
+        return raw_auth
+
+    if auth_type := properties.get(f"{AUTH}.type"):
+        type_prefix = f"{AUTH}.{auth_type}."
+        return {
+            "type": auth_type,
+            "impl": properties.get(f"{AUTH}.impl"),
+            auth_type: {
+                key[len(type_prefix) :].replace("-", "_"): value
+                for key, value in properties.items()
+                if key.startswith(type_prefix)
+            },
+        }
+
+    return None
+
+
+def _resolve_auth_config(auth_config: AuthConfigEnvelope | dict[str, Any]) -> ParsedAuthConfig:
+    auth_type = auth_config.get("type")
+    if not isinstance(auth_type, str):
+        raise ValueError("auth.type must be defined")
+
+    auth_type_config = auth_config.get(auth_type, {})
+    if not isinstance(auth_type_config, dict):
+        raise ValueError(f"auth.{auth_type} must be a dictionary")
+
+    auth_impl = auth_config.get("impl")
+    if auth_impl is not None and not isinstance(auth_impl, str):
+        raise ValueError("auth.impl must be a string")
+
+    auth_manager_name = auth_impl or auth_type
+
+    if auth_type == CUSTOM and not auth_impl:
+        raise ValueError("auth.impl must be specified when using custom auth.type")
+
+    if auth_type != CUSTOM and auth_impl:
+        raise ValueError("auth.impl can only be specified when using custom auth.type")
+
+    return {
+        "auth_type": auth_type,
+        "auth_manager_name": auth_manager_name,
+        "auth_type_config": auth_type_config,
+    }
+
+
 class RestCatalog(Catalog):
     uri: str
     _session: Session
@@ -500,47 +570,13 @@ def _create_session(self) -> Session:
                 elif ssl_client_cert := ssl_client.get(CERT):
                     session.cert = ssl_client_cert
 
-        raw_auth = self.properties.get(AUTH)
-        if isinstance(raw_auth, str):
-            try:
-                auth_config: dict[str, Any] | None = json.loads(raw_auth)
-            except json.JSONDecodeError as e:
-                raise ValueError("Failed to parse auth configuration as JSON") from e
-        elif raw_auth is not None:
-            auth_config = raw_auth
-        elif auth_type := self.properties.get(f"{AUTH}.type"):
-            type_prefix = f"{AUTH}.{auth_type}."
-            auth_config = {
-                "type": auth_type,
-                "impl": self.properties.get(f"{AUTH}.impl"),
-                auth_type: {
-                    key[len(type_prefix) :].replace("-", "_"): value
-                    for key, value in self.properties.items()
-                    if key.startswith(type_prefix)
-                },
-            }
-        else:
-            auth_config = None
-
-        if auth_config is not None and not isinstance(auth_config, dict):
-            raise ValueError("auth configuration must be a dictionary")
-
+        auth_config = _load_auth_config_from_properties(self.properties)
         if auth_config:
-            auth_type = auth_config.get("type")
-            if auth_type is None:
-                raise ValueError("auth.type must be defined")
-            auth_type_config = auth_config.get(auth_type, {})
-            auth_impl = auth_config.get("impl")
-            auth_manager_name = auth_impl or auth_type
-
-            if auth_type == CUSTOM and not auth_impl:
-                raise ValueError("auth.impl must be specified when using custom auth.type")
-
-            if auth_type != CUSTOM and auth_impl:
-                raise ValueError("auth.impl can only be specified when using custom auth.type")
-
-            typed_auth_type_config = _coerce_auth_config_values(auth_manager_name, auth_type_config)
-            self._auth_manager = AuthManagerFactory.create(auth_manager_name, typed_auth_type_config)
+            resolved_auth = _resolve_auth_config(auth_config)
+            typed_auth_type_config = _coerce_auth_config_values(
+                resolved_auth["auth_manager_name"], resolved_auth["auth_type_config"]
+            )
+            self._auth_manager = AuthManagerFactory.create(resolved_auth["auth_manager_name"], typed_auth_type_config)
             session.auth = AuthManagerAdapter(self._auth_manager)
         else:
             self._auth_manager = self._create_legacy_oauth2_auth_manager(session)