nucliadb-utils 5.0.0.post796__py3-none-any.whl → 5.0.0.post809__py3-none-any.whl

nucliadb_utils/aiopynecone/client.py

@@ -25,9 +25,15 @@ from collections.abc import AsyncIterable, Iterable
 from itertools import islice
 from typing import Any, AsyncGenerator, Optional
 
+import backoff
 import httpx
 
 from nucliadb_telemetry.metrics import Observer
+from nucliadb_utils.aiopynecone.exceptions import (
+    PineconeAPIError,
+    PineconeRateLimitError,
+    raise_for_status,
+)
 from nucliadb_utils.aiopynecone.models import (
     CreateIndexResponse,
     ListResponse,
@@ -38,9 +44,13 @@ from nucliadb_utils.aiopynecone.models import (
 
 logger = logging.getLogger(__name__)
 
+
 pinecone_observer = Observer(
     "pinecone_client",
     labels={"type": ""},
+    error_mappings={
+        "rate_limit": PineconeRateLimitError,
+    },
 )
 
 DEFAULT_TIMEOUT = 30
@@ -55,25 +65,18 @@ MAX_UPSERT_PAYLOAD_SIZE = 2 * MEGA_BYTE
 MAX_DELETE_BATCH_SIZE = 1000
 
 
-class PineconeAPIError(Exception):
-    def __init__(
-        self,
-        http_status_code: int,
-        code: Optional[str] = None,
-        message: Optional[str] = None,
-        details: Optional[Any] = None,
-    ):
-        self.http_status_code = http_status_code
-        self.code = code or ""
-        self.message = message or ""
-        self.details = details or {}
-        exc_message = '[{http_status_code}] message="{message}" code={code} details={details}'.format(
-            http_status_code=http_status_code,
-            message=message,
-            code=code,
-            details=details,
-        )
-        super().__init__(exc_message)
+RETRIABLE_EXCEPTIONS = (
+    PineconeRateLimitError,
+    httpx.ConnectError,
+    httpx.NetworkError,
+)
+
+backoff_handler = backoff.on_exception(
+    backoff.expo,
+    RETRIABLE_EXCEPTIONS,
+    jitter=backoff.random_jitter,
+    max_tries=4,
+)
 
 
 class ControlPlane:
@@ -105,7 +108,7 @@ class ControlPlane:
         }
         headers = {"Api-Key": self.api_key}
         http_response = await self.http_session.post("/indexes", json=payload, headers=headers)
-        raise_for_status(http_response)
+        raise_for_status("create_index", http_response)
         response = CreateIndexResponse.model_validate(http_response.json())
         return response.host
 
@@ -121,7 +124,7 @@ class ControlPlane:
         if response.status_code == 404:  # pragma: no cover
             logger.warning("Pinecone index not found.", extra={"index_name": name})
             return
-        raise_for_status(response)
+        raise_for_status("delete_index", response)
 
 
 class DataPlane:
@@ -147,6 +150,7 @@ class DataPlane:
     def _get_request_timeout(self, timeout: Optional[float] = None) -> Optional[float]:
         return timeout or self.client_timeout
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "upsert"})
     async def upsert(self, vectors: list[Vector], timeout: Optional[float] = None) -> None:
         """
@@ -165,7 +169,7 @@ class DataPlane:
         if request_timeout is not None:
             post_kwargs["timeout"] = timeout
         response = await self.http_session.post("/vectors/upsert", **post_kwargs)
-        raise_for_status(response)
+        raise_for_status("upsert", response)
 
     def _estimate_upsert_batch_size(self, vectors: list[Vector]) -> int:
         """
@@ -220,6 +224,7 @@ class DataPlane:
 
         await asyncio.gather(*tasks)
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "delete"})
     async def delete(self, ids: list[str], timeout: Optional[float] = None) -> None:
         """
@@ -242,8 +247,9 @@ class DataPlane:
         if request_timeout is not None:
             post_kwargs["timeout"] = timeout
         response = await self.http_session.post("/vectors/delete", **post_kwargs)
-        raise_for_status(response)
+        raise_for_status("delete", response)
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "list_page"})
     async def list_page(
         self,
@@ -279,7 +285,7 @@ class DataPlane:
             "/vectors/list",
             **post_kwargs,
         )
-        raise_for_status(response)
+        raise_for_status("list_page", response)
         return ListResponse.model_validate(response.json())
 
     async def list_all(
@@ -306,6 +312,7 @@ class DataPlane:
                 break
             pagination_token = response.pagination.next
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "delete_all"})
     async def delete_all(self, timeout: Optional[float] = None):
         """
@@ -324,7 +331,7 @@ class DataPlane:
             post_kwargs["timeout"] = timeout
         response = await self.http_session.post("/vectors/delete", **post_kwargs)
         try:
-            raise_for_status(response)
+            raise_for_status("delete_all", response)
         except PineconeAPIError as err:
             if err.http_status_code == 404 and err.code == 5:  # pragma: no cover
                 # Namespace not found. No vectors to delete.
@@ -366,6 +373,7 @@ class DataPlane:
 
         await asyncio.gather(*tasks)
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "query"})
     async def query(
         self,
@@ -404,7 +412,7 @@ class DataPlane:
         if request_timeout is not None:
             post_kwargs["timeout"] = timeout
         response = await self.http_session.post("/query", **post_kwargs)
-        raise_for_status(response)
+        raise_for_status("query", response)
         return QueryResponse.model_validate(response.json())
 
 
@@ -459,28 +467,6 @@ class PineconeSession:
         return DataPlane(api_key=api_key, index_host_session=index_host_session, timeout=timeout)
 
 
-def raise_for_status(response: httpx.Response):
-    try:
-        response.raise_for_status()
-    except httpx.HTTPStatusError:
-        code = None
-        message = None
-        details = None
-        try:
-            resp_json = response.json()
-            code = resp_json.get("code")
-            message = resp_json.get("message")
-            details = resp_json.get("details")
-        except Exception:
-            message = response.text
-        raise PineconeAPIError(
-            http_status_code=response.status_code,
-            code=code,
-            message=message,
-            details=details,
-        )
-
-
 def batchify(iterable: Iterable, batch_size: int):
     """
     Split an iterable into batches of batch_size

nucliadb_utils/aiopynecone/exceptions.py

@@ -0,0 +1,90 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 Any, Optional
+
+import httpx
+
+from nucliadb_telemetry.metrics import Counter
+
+pinecone_errors_counter = Counter("pinecone_errors", labels={"type": ""})
+
+
+class PineconeAPIError(Exception):
+    """
+    Generic Pinecone API error.
+    """
+
+    def __init__(
+        self,
+        http_status_code: int,
+        code: Optional[str] = None,
+        message: Optional[str] = None,
+        details: Optional[Any] = None,
+    ):
+        self.http_status_code = http_status_code
+        self.code = code or ""
+        self.message = message or ""
+        self.details = details or {}
+        exc_message = '[{http_status_code}] message="{message}" code={code} details={details}'.format(
+            http_status_code=http_status_code,
+            message=message,
+            code=code,
+            details=details,
+        )
+        super().__init__(exc_message)
+
+
+class PineconeRateLimitError(PineconeAPIError):
+    """
+    Raised when the client has exceeded the rate limit to be able to backoff and retry.
+    """
+
+    pass
+
+
+def raise_for_status(operation: str, response: httpx.Response):
+    try:
+        response.raise_for_status()
+    except httpx.HTTPStatusError:
+        pinecone_errors_counter.inc(labels={"type": operation})
+        code = None
+        message = None
+        details = None
+        try:
+            resp_json = response.json()
+            code = resp_json.get("code")
+            message = resp_json.get("message")
+            details = resp_json.get("details")
+        except Exception:
+            message = response.text
+        if response.status_code == 429:
+            raise PineconeRateLimitError(
+                http_status_code=response.status_code,
+                code=code,
+                message=message,
+                details=details,
+            )
+        raise PineconeAPIError(
+            http_status_code=response.status_code,
+            code=code,
+            message=message,
+            details=details,
+        )

{nucliadb_utils-5.0.0.post796.dist-info → nucliadb_utils-5.0.0.post809.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb_utils
-Version: 5.0.0.post796
+Version: 5.0.0.post809
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -23,8 +23,8 @@ Requires-Dist: PyNaCl
 Requires-Dist: pyjwt >=2.4.0
 Requires-Dist: memorylru >=1.1.2
 Requires-Dist: mrflagly
-Requires-Dist: nucliadb-protos >=5.0.0.post796
-Requires-Dist: nucliadb-telemetry >=5.0.0.post796
+Requires-Dist: nucliadb-protos >=5.0.0.post809
+Requires-Dist: nucliadb-telemetry >=5.0.0.post809
 Provides-Extra: cache
 Requires-Dist: redis >=4.3.4 ; extra == 'cache'
 Requires-Dist: orjson >=3.6.7 ; extra == 'cache'

{nucliadb_utils-5.0.0.post796.dist-info → nucliadb_utils-5.0.0.post809.dist-info}/RECORD

@@ -18,7 +18,8 @@ nucliadb_utils/store.py,sha256=kQ35HemE0v4_Qg6xVqNIJi8vSFAYQtwI3rDtMsNy62Y,890
 nucliadb_utils/transaction.py,sha256=mwcI3aIHAvU5KOGqd_Uz_d1XQzXhk_-NWY8NqU1lfb0,7307
 nucliadb_utils/utilities.py,sha256=oz3tEODG2g3todnyvA-nW1Ou6xXDveL_tMKTDGdWXM4,15287
 nucliadb_utils/aiopynecone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb_utils/aiopynecone/client.py,sha256=kvGLCzSbDaMIWx0LK9WBAL7QsyuPmUjC7cPv5djMdFw,19028
+nucliadb_utils/aiopynecone/client.py,sha256=T0S6aoVefnjyKgJAGzBJGnrhBDlvSAj5DJBqdGRuzZw,18407
+nucliadb_utils/aiopynecone/exceptions.py,sha256=JJ-Ui9TSoHRe8WfkHR04S-t7NsPCy1xxluPH5R7fDl8,2816
 nucliadb_utils/aiopynecone/models.py,sha256=DVlCVrinHAaDxuii3fzdciFn4dqS5HTTyBavnoDlR2U,2024
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/audit/audit.py,sha256=dn5ZnCVQUlCcvdjzaORghbrjk9QgVGrtkfIftq30Bp8,2819
@@ -63,8 +64,8 @@ nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQ
 nucliadb_utils/tests/local.py,sha256=c3gZJJWmvOftruJkIQIwB3q_hh3uxEhqGIAVWim1Bbk,1343
 nucliadb_utils/tests/nats.py,sha256=Tosonm9A9cusImyji80G4pgdXEHNVPaCLT5TbFK_ra0,7543
 nucliadb_utils/tests/s3.py,sha256=YB8QqDaBXxyhHonEHmeBbRRDmvB7sTOaKBSi8KBGokg,2330
-nucliadb_utils-5.0.0.post796.dist-info/METADATA,sha256=mwdR7jpGf8oA1T8pa62MXf_BajMORpePKeK7KpiujIM,2073
-nucliadb_utils-5.0.0.post796.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-nucliadb_utils-5.0.0.post796.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-5.0.0.post796.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-5.0.0.post796.dist-info/RECORD,,
+nucliadb_utils-5.0.0.post809.dist-info/METADATA,sha256=j1oedwxyH4rR_jpUuyUgFWnAcNImYZtGXSGvaIiXtnE,2073
+nucliadb_utils-5.0.0.post809.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+nucliadb_utils-5.0.0.post809.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-5.0.0.post809.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-5.0.0.post809.dist-info/RECORD,,