nucliadb-utils 5.0.1.post1091__py3-none-any.whl → 5.0.1.post1101__py3-none-any.whl

nucliadb_utils/aiopynecone/client.py

@@ -28,7 +28,7 @@ from typing import Any, AsyncGenerator, Optional
 import backoff
 import httpx
 
-from nucliadb_telemetry.metrics import Observer
+from nucliadb_telemetry.metrics import INF, Histogram, Observer
 from nucliadb_utils.aiopynecone.exceptions import (
     PineconeAPIError,
     PineconeRateLimitError,
@@ -46,6 +46,25 @@ from nucliadb_utils.aiopynecone.models import (
 
 logger = logging.getLogger(__name__)
 
+upsert_batch_size_histogram = Histogram(
+    "pinecone_upsert_batch_size",
+    buckets=[10.0, 100.0, 200.0, 500.0, 1000.0, 5000.0, INF],
+)
+upsert_batch_count_histogram = Histogram(
+    "pinecone_upsert_batch_count",
+    buckets=[0.0, 1.0, 2.0, 3.0, 5.0, 10.0, 15.0, 20.0, 30.0, 50.0, INF],
+)
+
+delete_batch_size_histogram = Histogram(
+    "pinecone_delete_batch_size",
+    buckets=[1.0, 5.0, 10.0, 20.0, 50.0, 100.0, 150.0, INF],
+)
+
+delete_batch_count_histogram = Histogram(
+    "pinecone_delete_batch_count",
+    buckets=[0.0, 1.0, 2.0, 3.0, 5.0, 10.0, 15.0, 20.0, 30.0, 50.0, INF],
+)
+
 
 pinecone_observer = Observer(
     "pinecone_client",
@@ -57,7 +76,6 @@ pinecone_observer = Observer(
 
 DEFAULT_TIMEOUT = 30
 CONTROL_PLANE_BASE_URL = "https://api.pinecone.io/"
-INDEX_HOST_BASE_URL = "https://{index_host}/"
 BASE_API_HEADERS = {
     "Content-Type": "application/json",
     "Accept": "application/json",
@@ -198,6 +216,7 @@ class DataPlane:
         if len(vectors) == 0:
             # Nothing to upsert.
             return
+        upsert_batch_size_histogram.observe(len(vectors))
         headers = {"Api-Key": self.api_key}
         payload = UpsertRequest(vectors=vectors)
         post_kwargs: dict[str, Any] = {
@@ -261,7 +280,10 @@ class DataPlane:
         for batch in batchify(vectors, batch_size):
             tasks.append(asyncio.create_task(_upsert_batch(batch)))
 
-        await asyncio.gather(*tasks)
+        upsert_batch_count_histogram.observe(len(tasks))
+
+        if len(tasks) > 0:
+            await asyncio.gather(*tasks)
 
     @backoff.on_exception(
         backoff.expo,
@@ -280,8 +302,16 @@ class DataPlane:
         """
         if len(ids) > MAX_DELETE_BATCH_SIZE:
             raise ValueError(f"Maximum number of ids in a single request is {MAX_DELETE_BATCH_SIZE}.")
+        if len(ids) == 0:  # pragma: no cover
+            return
 
+        delete_batch_size_histogram.observe(len(ids))
         headers = {"Api-Key": self.api_key}
+
+        # This is a temporary log info to hunt down a bug.
+        rids = {vid.split("/")[0] for vid in ids}
+        logger.info(f"Deleting vectors from resources: {list(rids)}")
+
         payload = {"ids": ids}
         post_kwargs: dict[str, Any] = {
             "headers": headers,
@@ -428,7 +458,10 @@ class DataPlane:
         async for batch in async_batchify(async_iterable, batch_size):
             tasks.append(asyncio.create_task(_delete_batch(batch)))
 
-        await asyncio.gather(*tasks)
+        delete_batch_count_histogram.observe(len(tasks))
+
+        if len(tasks) > 0:
+            await asyncio.gather(*tasks)
 
     @backoff.on_exception(
         backoff.expo,
@@ -516,8 +549,12 @@ class PineconeSession:
         if session is not None:
             return session
 
+        base_url = index_host
+        if not index_host.startswith("https://"):
+            base_url = f"https://{index_host}/"
+
         session = httpx.AsyncClient(
-            base_url=INDEX_HOST_BASE_URL.format(index_host=index_host),
+            base_url=base_url,
             headers=BASE_API_HEADERS,
             timeout=DEFAULT_TIMEOUT,
         )

nucliadb_utils/storages/gcs.py

@@ -34,6 +34,7 @@ import aiohttp.client_exceptions
 import backoff
 import google.auth.transport.requests  # type: ignore
 import yarl
+from google.auth.exceptions import DefaultCredentialsError  # type: ignore
 from google.oauth2 import service_account  # type: ignore
 
 from nucliadb_protos.resources_pb2 import CloudFile
@@ -458,12 +459,25 @@ class GCSStorage(Storage):
         url: str = "https://www.googleapis.com",
         scopes: Optional[List[str]] = None,
     ):
-        if account_credentials is not None:
+        if account_credentials is None:
+            self._json_credentials = None
+        elif isinstance(account_credentials, str) and account_credentials.strip() == "":
+            self._json_credentials = None
+        else:
             self._json_credentials = json.loads(base64.b64decode(account_credentials))
+
+        if self._json_credentials is not None:
             self._credentials = service_account.Credentials.from_service_account_info(
                 self._json_credentials,
                 scopes=DEFAULT_SCOPES if scopes is None else scopes,
             )
+        else:
+            try:
+                self._credentials, self._project = google.auth.default()
+            except DefaultCredentialsError:
+                logger.warning("Setting up without credentials as couldn't find workload identity")
+                self._credentials = None
+
         self.source = CloudFile.GCS
         self.deadletter_bucket = deadletter_bucket
         self.indexing_bucket = indexing_bucket
@@ -473,16 +487,15 @@ class GCSStorage(Storage):
         # https://cloud.google.com/storage/docs/bucket-locations
         self._bucket_labels = labels or {}
         self._executor = executor
-        self._creation_access_token = datetime.now()
         self._upload_url = url + "/upload/storage/v1/b/{bucket}/o?uploadType=resumable"  # noqa
         self.object_base_url = url + "/storage/v1/b"
         self._client = None
 
     def _get_access_token(self):
-        if self._credentials.valid is False:
-            req = google.auth.transport.requests.Request()
-            self._credentials.refresh(req)
-            self._creation_access_token = datetime.now()
+        if self._credentials.expired or self._credentials.valid is False:
+            request = google.auth.transport.requests.Request()
+            self._credentials.refresh(request)
+
         return self._credentials.token
 
     @storage_ops_observer.wrap({"type": "initialize"})
@@ -552,7 +565,9 @@ class GCSStorage(Storage):
             raise AttributeError()
 
         headers = await self.get_access_headers()
-        url = f"{self.object_base_url}/{bucket_name}?project={self._project}"
+        # Using object access url instead of bucket access to avoid
+        # giving admin permission to the SA, needed to GET a bucket
+        url = f"{self.object_base_url}/{bucket_name}/o"
         async with self.session.get(
             url,
             headers=headers,

{nucliadb_utils-5.0.1.post1091.dist-info → nucliadb_utils-5.0.1.post1101.dist-info}/METADATA

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

{nucliadb_utils-5.0.1.post1091.dist-info → nucliadb_utils-5.0.1.post1101.dist-info}/RECORD

@@ -18,7 +18,7 @@ nucliadb_utils/store.py,sha256=kQ35HemE0v4_Qg6xVqNIJi8vSFAYQtwI3rDtMsNy62Y,890
 nucliadb_utils/transaction.py,sha256=mwcI3aIHAvU5KOGqd_Uz_d1XQzXhk_-NWY8NqU1lfb0,7307
 nucliadb_utils/utilities.py,sha256=idajCm_4Sojh7b3HTkP0fTfG2Mb6PIB9xtMmcfB7Nl0,15758
 nucliadb_utils/aiopynecone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb_utils/aiopynecone/client.py,sha256=EFAQbNp-OHZ55vyXWzdzzp_59FGacGs_nfM2sMBIG8A,21153
+nucliadb_utils/aiopynecone/client.py,sha256=3DGJuHXO9Vm83atMLD4fd2FTz8SAr8RvPCGU_bXH9Ho,22333
 nucliadb_utils/aiopynecone/exceptions.py,sha256=EEE0XoGs1zIB5yOJ_fy6yoG4uIb4cWIawYdJeNe4eDo,3012
 nucliadb_utils/aiopynecone/models.py,sha256=ketK2IYLWiwFZ76rnJmwfcuopFJrCAtCUszdTSurm_Q,3236
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
@@ -47,7 +47,7 @@ nucliadb_utils/nuclia_usage/utils/kb_usage_report.py,sha256=lTr9CMBpdk34KtkH5K8v
 nucliadb_utils/storages/__init__.py,sha256=5Qc8AUWiJv9_JbGCBpAn88AIJhwDlm0OPQpg2ZdRL4U,872
 nucliadb_utils/storages/azure.py,sha256=egMDwLNIGSQyVevuySt2AswzFdNAcih05BbRg3-p8IU,16015
 nucliadb_utils/storages/exceptions.py,sha256=mm_wX4YRtp7u7enkk_4pMSlX5AQQuFbq4xLmupVDt3Y,2502
-nucliadb_utils/storages/gcs.py,sha256=k1WXZDBQXSizO_3kEJef7EmyInnASJsAy5VTtEw8ly4,27020
+nucliadb_utils/storages/gcs.py,sha256=TQlPpg9HzfHbOu_iCLPUo9KAql1IIjWxJ9mVJ9_CrV8,27658
 nucliadb_utils/storages/local.py,sha256=NxC_nMBd38NDsR266DSgoBLdQlvUwf0_sd50r-BLI0E,10288
 nucliadb_utils/storages/nuclia.py,sha256=vEv94xAT7QM2g80S25QyrOw2pzvP2BAX-ADgZLtuCVc,2097
 nucliadb_utils/storages/object_store.py,sha256=Tw10GmpYfM5TMqJ3Tk9pLQ9wLMBk1-snL_m6uasiZDQ,4257
@@ -64,8 +64,8 @@ nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQ
 nucliadb_utils/tests/local.py,sha256=7nuP8EFUAiA8ZH50R1iPV9EUXBySQxOanVm3Zht_e0g,1835
 nucliadb_utils/tests/nats.py,sha256=xqpww4jZjTKY9oPGlJdDJG67L3FIBQsa9qDHxILR8r8,7687
 nucliadb_utils/tests/s3.py,sha256=IdMxK_cNdSHLvO1u8BwsKFzD87Hk1MVPDZ57zx6h-rA,3656
-nucliadb_utils-5.0.1.post1091.dist-info/METADATA,sha256=36ObBdcAROkELutBroI5F89yEBH8CxZ5B2Gj90RA9_c,2071
-nucliadb_utils-5.0.1.post1091.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-nucliadb_utils-5.0.1.post1091.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-5.0.1.post1091.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-5.0.1.post1091.dist-info/RECORD,,
+nucliadb_utils-5.0.1.post1101.dist-info/METADATA,sha256=QPBsHrTmcgVbSx8pao1HXnknyhDjP0j8TWCjRvzYcOA,2071
+nucliadb_utils-5.0.1.post1101.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+nucliadb_utils-5.0.1.post1101.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-5.0.1.post1101.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-5.0.1.post1101.dist-info/RECORD,,