This is an automated email from the ASF dual-hosted git repository.
honahx pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
The following commit(s) were added to refs/heads/main by this push:
new 1f433a46 Fix retrying logic (#480)
1f433a46 is described below
commit 1f433a46ca59a896014369acfd8187bed4ca4fa2
Author: Fokko Driesprong <[email protected]>
AuthorDate: Thu Feb 29 09:31:29 2024 +0100
Fix retrying logic (#480)
---
pyiceberg/catalog/rest.py | 8 +++----
tests/catalog/test_rest.py | 58 +++++++++++++++++++++++++++++++---------------
2 files changed, 43 insertions(+), 23 deletions(-)
diff --git a/pyiceberg/catalog/rest.py b/pyiceberg/catalog/rest.py
index 6b952ef0..3132af9f 100644
--- a/pyiceberg/catalog/rest.py
+++ b/pyiceberg/catalog/rest.py
@@ -128,7 +128,7 @@ def _retry_hook(retry_state: RetryCallState) -> None:
_RETRY_ARGS = {
"retry": retry_if_exception_type(AuthorizationExpiredError),
"stop": stop_after_attempt(2),
- "before": _retry_hook,
+ "before_sleep": _retry_hook,
"reraise": True,
}
@@ -446,10 +446,10 @@ class RestCatalog(Catalog):
catalog=self,
)
- def _refresh_token(self, session: Optional[Session] = None, new_token:
Optional[str] = None) -> None:
+ def _refresh_token(self, session: Optional[Session] = None, initial_token:
Optional[str] = None) -> None:
session = session or self._session
- if new_token is not None:
- self.properties[TOKEN] = new_token
+ if initial_token is not None:
+ self.properties[TOKEN] = initial_token
elif CREDENTIAL in self.properties:
self.properties[TOKEN] = self._fetch_access_token(session,
self.properties[CREDENTIAL])
diff --git a/tests/catalog/test_rest.py b/tests/catalog/test_rest.py
index fa35ec45..21b4d955 100644
--- a/tests/catalog/test_rest.py
+++ b/tests/catalog/test_rest.py
@@ -306,24 +306,39 @@ def test_list_namespace_with_parent_200(rest_mock:
Mocker) -> None:
]
-def test_list_namespaces_419(rest_mock: Mocker) -> None:
+def test_list_namespaces_token_expired(rest_mock: Mocker) -> None:
new_token = "new_jwt_token"
new_header = dict(TEST_HEADERS)
new_header["Authorization"] = f"Bearer {new_token}"
- rest_mock.post(
+ namespaces = rest_mock.register_uri(
+ "GET",
f"{TEST_URI}v1/namespaces",
- json={
- "error": {
- "message": "Authorization expired.",
- "type": "AuthorizationExpiredError",
- "code": 419,
- }
- },
- status_code=419,
- request_headers=TEST_HEADERS,
- )
- rest_mock.post(
+ [
+ {
+ "status_code": 419,
+ "json": {
+ "error": {
+ "message": "Authorization expired.",
+ "type": "AuthorizationExpiredError",
+ "code": 419,
+ }
+ },
+ "headers": TEST_HEADERS,
+ },
+ {
+ "status_code": 200,
+ "json": {"namespaces": [["default"], ["examples"], ["fokko"],
["system"]]},
+ "headers": new_header,
+ },
+ {
+ "status_code": 200,
+ "json": {"namespaces": [["default"], ["examples"], ["fokko"],
["system"]]},
+ "headers": new_header,
+ },
+ ],
+ )
+ tokens = rest_mock.post(
f"{TEST_URI}v1/oauth/tokens",
json={
"access_token": new_token,
@@ -333,12 +348,6 @@ def test_list_namespaces_419(rest_mock: Mocker) -> None:
},
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",),
@@ -346,6 +355,17 @@ def test_list_namespaces_419(rest_mock: Mocker) -> None:
("fokko",),
("system",),
]
+ assert namespaces.call_count == 2
+ assert tokens.call_count == 1
+
+ assert catalog.list_namespaces() == [
+ ("default",),
+ ("examples",),
+ ("fokko",),
+ ("system",),
+ ]
+ assert namespaces.call_count == 3
+ assert tokens.call_count == 1
def test_create_namespace_200(rest_mock: Mocker) -> None:
