Retry with new Access Token HTTP 419 (#340)

* Refresh Auth token on expiry

* Check call count

* Add test to cover retry logic

* Update poetry.lock with tenacity

* Fix tests for Python <= 3.9

Author: anupam-saini

poetry.lock

@@ -30,6 +30,7 @@
 
 from pydantic import Field, ValidationError
 from requests import HTTPError, Session
+from tenacity import RetryCallState, retry, retry_if_exception_type, stop_after_attempt
 
 from pyiceberg import __version__
 from pyiceberg.catalog import (
@@ -118,6 +119,19 @@ class Endpoints:
 NAMESPACE_SEPARATOR = b"\x1F".decode(UTF8)
 
 
+def _retry_hook(retry_state: RetryCallState) -> None:
+    rest_catalog: RestCatalog = retry_state.args[0]
+    rest_catalog._refresh_token()  # pylint: disable=protected-access
+
+
+_RETRY_ARGS = {
+    "retry": retry_if_exception_type(AuthorizationExpiredError),
+    "stop": stop_after_attempt(2),
+    "before": _retry_hook,
+    "reraise": True,
+}
+
+
 class TableResponse(IcebergBaseModel):
     metadata_location: str = Field(alias="metadata-location")
     metadata: TableMetadata
@@ -225,13 +239,7 @@ def _create_session(self) -> Session:
                 elif ssl_client_cert := ssl_client.get(CERT):
                     session.cert = ssl_client_cert
 
-        # If we have credentials, but not a token, we want to fetch a token
-        if TOKEN not in self.properties and CREDENTIAL in self.properties:
-            self.properties[TOKEN] = self._fetch_access_token(session, self.properties[CREDENTIAL])
-
-        # Set Auth token for subsequent calls in the session
-        if token := self.properties.get(TOKEN):
-            session.headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {token}"
+        self._refresh_token(session, self.properties.get(TOKEN))
 
         # Set HTTP headers
         session.headers["Content-type"] = "application/json"
@@ -439,6 +447,18 @@ def _response_to_table(self, identifier_tuple: Tuple[str, ...], table_response:
             catalog=self,
         )
 
+    def _refresh_token(self, session: Optional[Session] = None, new_token: Optional[str] = None) -> None:
+        session = session or self._session
+        if new_token is not None:
+            self.properties[TOKEN] = new_token
+        elif CREDENTIAL in self.properties:
+            self.properties[TOKEN] = self._fetch_access_token(session, self.properties[CREDENTIAL])
+
+        # Set Auth token for subsequent calls in the session
+        if token := self.properties.get(TOKEN):
+            session.headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {token}"
+
+    @retry(**_RETRY_ARGS)
     def create_table(
         self,
         identifier: Union[str, Identifier],
@@ -475,6 +495,7 @@ def create_table(
         table_response = TableResponse(**response.json())
         return self._response_to_table(self.identifier_to_tuple(identifier), table_response)
 
+    @retry(**_RETRY_ARGS)
     def register_table(self, identifier: Union[str, Identifier], metadata_location: str) -> Table:
         """Register a new table using existing metadata.
 
@@ -506,6 +527,7 @@ def register_table(self, identifier: Union[str, Identifier], metadata_location:
         table_response = TableResponse(**response.json())
         return self._response_to_table(self.identifier_to_tuple(identifier), table_response)
 
+    @retry(**_RETRY_ARGS)
     def list_tables(self, namespace: Union[str, Identifier]) -> List[Identifier]:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace_concat = NAMESPACE_SEPARATOR.join(namespace_tuple)
@@ -516,6 +538,7 @@ def list_tables(self, namespace: Union[str, Identifier]) -> List[Identifier]:
             self._handle_non_200_response(exc, {404: NoSuchNamespaceError})
         return [(*table.namespace, table.name) for table in ListTablesResponse(**response.json()).identifiers]
 
+    @retry(**_RETRY_ARGS)
     def load_table(self, identifier: Union[str, Identifier]) -> Table:
         identifier_tuple = self.identifier_to_tuple_without_catalog(identifier)
         response = self._session.get(
@@ -529,6 +552,7 @@ def load_table(self, identifier: Union[str, Identifier]) -> Table:
         table_response = TableResponse(**response.json())
         return self._response_to_table(identifier_tuple, table_response)
 
+    @retry(**_RETRY_ARGS)
     def drop_table(self, identifier: Union[str, Identifier], purge_requested: bool = False) -> None:
         identifier_tuple = self.identifier_to_tuple_without_catalog(identifier)
         response = self._session.delete(
@@ -541,9 +565,11 @@ def drop_table(self, identifier: Union[str, Identifier], purge_requested: bool =
         except HTTPError as exc:
             self._handle_non_200_response(exc, {404: NoSuchTableError})
 
+    @retry(**_RETRY_ARGS)
     def purge_table(self, identifier: Union[str, Identifier]) -> None:
         self.drop_table(identifier=identifier, purge_requested=True)
 
+    @retry(**_RETRY_ARGS)
     def rename_table(self, from_identifier: Union[str, Identifier], to_identifier: Union[str, Identifier]) -> Table:
         from_identifier_tuple = self.identifier_to_tuple_without_catalog(from_identifier)
         payload = {
@@ -558,6 +584,7 @@ def rename_table(self, from_identifier: Union[str, Identifier], to_identifier: U
 
         return self.load_table(to_identifier)
 
+    @retry(**_RETRY_ARGS)
     def _commit_table(self, table_request: CommitTableRequest) -> CommitTableResponse:
         """Update the table.
 
@@ -588,6 +615,7 @@ def _commit_table(self, table_request: CommitTableRequest) -> CommitTableRespons
             )
         return CommitTableResponse(**response.json())
 
+    @retry(**_RETRY_ARGS)
     def create_namespace(self, namespace: Union[str, Identifier], properties: Properties = EMPTY_DICT) -> None:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         payload = {"namespace": namespace_tuple, "properties": properties}
@@ -597,6 +625,7 @@ def create_namespace(self, namespace: Union[str, Identifier], properties: Proper
         except HTTPError as exc:
             self._handle_non_200_response(exc, {404: NoSuchNamespaceError, 409: NamespaceAlreadyExistsError})
 
+    @retry(**_RETRY_ARGS)
     def drop_namespace(self, namespace: Union[str, Identifier]) -> None:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace = NAMESPACE_SEPARATOR.join(namespace_tuple)
@@ -606,6 +635,7 @@ def drop_namespace(self, namespace: Union[str, Identifier]) -> None:
         except HTTPError as exc:
             self._handle_non_200_response(exc, {404: NoSuchNamespaceError})
 
+    @retry(**_RETRY_ARGS)
     def list_namespaces(self, namespace: Union[str, Identifier] = ()) -> List[Identifier]:
         namespace_tuple = self.identifier_to_tuple(namespace)
         response = self._session.get(
@@ -623,6 +653,7 @@ def list_namespaces(self, namespace: Union[str, Identifier] = ()) -> List[Identi
         namespaces = ListNamespaceResponse(**response.json())
         return [namespace_tuple + child_namespace for child_namespace in namespaces.namespaces]
 
+    @retry(**_RETRY_ARGS)
     def load_namespace_properties(self, namespace: Union[str, Identifier]) -> Properties:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace = NAMESPACE_SEPARATOR.join(namespace_tuple)
@@ -634,6 +665,7 @@ def load_namespace_properties(self, namespace: Union[str, Identifier]) -> Proper
 
         return NamespaceResponse(**response.json()).properties
 
+    @retry(**_RETRY_ARGS)
     def update_namespace_properties(
         self, namespace: Union[str, Identifier], removals: Optional[Set[str]] = None, updates: Properties = EMPTY_DICT
     ) -> PropertiesUpdateSummary:

pyiceberg/catalog/rest.py

@@ -57,6 +57,7 @@ sortedcontainers = "2.4.0"
 fsspec = ">=2023.1.0,<2024.1.0"
 pyparsing = ">=3.1.0,<4.0.0"
 zstandard = ">=0.13.0,<1.0.0"
+tenacity = ">=8.2.3,<9.0.0"
 pyarrow = { version = ">=9.0.0,<16.0.0", optional = true }
 pandas = { version = ">=1.0.0,<3.0.0", optional = true }
 duckdb = { version = ">=0.5.0,<1.0.0", optional = true }
@@ -301,6 +302,10 @@ ignore_missing_imports = true
 module = "setuptools.*"
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = "tenacity.*"
+ignore_missing_imports = true
+
 [tool.coverage.run]
 source = ['pyiceberg/']
 

pyproject.toml

@@ -26,6 +26,7 @@
 from pyiceberg.catalog import PropertiesUpdateSummary, Table, load_catalog
 from pyiceberg.catalog.rest import AUTH_URL, RestCatalog
 from pyiceberg.exceptions import (
+    AuthorizationExpiredError,
     NamespaceAlreadyExistsError,
     NoSuchNamespaceError,
     NoSuchTableError,
@@ -266,6 +267,48 @@ def test_list_namespace_with_parent_200(rest_mock: Mocker) -> None:
     ]
 
 
+def test_list_namespaces_419(rest_mock: Mocker) -> None:
+    new_token = "new_jwt_token"
+    new_header = dict(TEST_HEADERS)
+    new_header["Authorization"] = f"Bearer {new_token}"
+
+    rest_mock.post(
+        f"{TEST_URI}v1/namespaces",
+        json={
+            "error": {
+                "message": "Authorization expired.",
+                "type": "AuthorizationExpiredError",
+                "code": 419,
+            }
+        },
+        status_code=419,
+        request_headers=TEST_HEADERS,
+    )
+    rest_mock.post(
+        f"{TEST_URI}v1/oauth/tokens",
+        json={
+            "access_token": new_token,
+            "token_type": "Bearer",
+            "expires_in": 86400,
+            "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
+        },
+        status_code=200,
+    )
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces",
+        json={"namespaces": [["default"], ["examples"], ["fokko"], ["system"]]},
+        status_code=200,
+        request_headers=new_header,
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, credential=TEST_CREDENTIALS)
+    assert catalog.list_namespaces() == [
+        ("default",),
+        ("examples",),
+        ("fokko",),
+        ("system",),
+    ]
+
+
 def test_create_namespace_200(rest_mock: Mocker) -> None:
     namespace = "leden"
     rest_mock.post(
@@ -517,6 +560,35 @@ def test_create_table_409(rest_mock: Mocker, table_schema_simple: Schema) -> Non
     assert "Table already exists" in str(e.value)
 
 
+def test_create_table_419(rest_mock: Mocker, table_schema_simple: Schema) -> None:
+    rest_mock.post(
+        f"{TEST_URI}v1/namespaces/fokko/tables",
+        json={
+            "error": {
+                "message": "Authorization expired.",
+                "type": "AuthorizationExpiredError",
+                "code": 419,
+            }
+        },
+        status_code=419,
+        request_headers=TEST_HEADERS,
+    )
+
+    with pytest.raises(AuthorizationExpiredError) as e:
+        RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN).create_table(
+            identifier=("fokko", "fokko2"),
+            schema=table_schema_simple,
+            location=None,
+            partition_spec=PartitionSpec(
+                PartitionField(source_id=1, field_id=1000, transform=TruncateTransform(width=3), name="id")
+            ),
+            sort_order=SortOrder(SortField(source_id=2, transform=IdentityTransform())),
+            properties={"owner": "fokko"},
+        )
+    assert "Authorization expired" in str(e.value)
+    assert rest_mock.call_count == 3
+
+
 def test_register_table_200(
     rest_mock: Mocker, table_schema_simple: Schema, example_table_metadata_no_snapshot_v1_rest_json: Dict[str, Any]
 ) -> None: