propagate try after in rate limit errors on sdk (#2461)

nucliadb_sdk/src/nucliadb_sdk/v2/exceptions.py

@@ -16,6 +16,9 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
+from typing import Optional
+
+
 class ClientError(Exception):
     pass
 
@@ -33,7 +36,9 @@ class AccountLimitError(ClientError):
 
 
 class RateLimitError(ClientError):
-    pass
+    def __init__(self, message, try_after: Optional[float] = None):
+        super().__init__(message)
+        self.try_after = try_after
 
 
 class ConflictError(ClientError):

nucliadb_sdk/src/nucliadb_sdk/v2/sdk.py

@@ -21,6 +21,7 @@
 import inspect
 import io
 import warnings
+from json import JSONDecodeError
 from typing import (
     Any,
     AsyncGenerator,
@@ -345,7 +346,13 @@ def _check_response(self, response: httpx.Response):
                 f"Account limits exceeded error {response.status_code}: {response.text}"
             )
         elif response.status_code == 429:
-            raise exceptions.RateLimitError(response.text)
+            try_after: Optional[float] = None
+            try:
+                body = response.json()
+                try_after = body.get("detail", {}).get("try_after")
+            except JSONDecodeError:
+                pass
+            raise exceptions.RateLimitError(response.text, try_after=try_after)
         elif response.status_code in (
             409,
             419,

nucliadb_sdk/tests/test_sdk.py

@@ -17,6 +17,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
+import httpx
 import pytest
 
 import nucliadb_sdk
@@ -103,3 +104,44 @@ def test_learning_config_endpoints(sdk: nucliadb_sdk.NucliaDB, kb):
     sdk.get_models(kbid=kb.uuid)
     sdk.get_model(kbid=kb.uuid, model_id="foo")
     sdk.get_configuration_schema(kbid=kb.uuid)
+
+
+def test_check_response():
+    sdk = nucliadb_sdk.NucliaDB(region="europe-1")
+
+    response = httpx.Response(200)
+    assert sdk._check_response(response) is response
+
+    response = httpx.Response(299)
+    assert sdk._check_response(response) is response
+
+    with pytest.raises(nucliadb_sdk.exceptions.UnknownError) as err:
+        sdk._check_response(httpx.Response(300, text="foo"))
+        assert str(err.value) == "Unknown error connecting to API: 300: foo"
+
+    for status_code in (401, 403):
+        with pytest.raises(nucliadb_sdk.exceptions.AuthError) as err:
+            sdk._check_response(httpx.Response(status_code, text="foo"))
+            assert str(err.value) == f"Auth error {status_code}: foo"
+
+    with pytest.raises(nucliadb_sdk.exceptions.AccountLimitError) as err:
+        sdk._check_response(httpx.Response(402, text="foo"))
+        assert str(err.value) == f"Account limits exceeded error {status_code}: foo"
+
+    with pytest.raises(nucliadb_sdk.exceptions.RateLimitError) as err:
+        sdk._check_response(httpx.Response(429, json={"detail": {"try_after": 1}}, text="Rate limit!"))
+        assert str(err.value) == f"Rate limit!"
+        assert err.value.try_after == 1
+
+    for status_code in (409, 419):
+        with pytest.raises(nucliadb_sdk.exceptions.ConflictError) as err:
+            sdk._check_response(httpx.Response(status_code, text="foo"))
+            assert str(err.value) == "foo"
+
+    with pytest.raises(nucliadb_sdk.exceptions.NotFoundError) as err:
+        sdk._check_response(
+            httpx.Response(
+                404, text="foo", request=httpx.Request(method="GET", url=httpx.URL("http://url"))
+            ),
+        )
+        assert str(err.value) == "Resource not found at http://url: foo"