fix(clients): allow chunked requests on WithTransformation methods (generated)

algolia/api-clients-automation#5011

Co-authored-by: algolia-bot <[email protected]>
Co-authored-by: Clément Vannicatte <[email protected]>

Author: algolia-bot

algoliasearch/ingestion/client.py

@@ -22,13 +22,20 @@
 
 from algoliasearch.http.api_response import ApiResponse
 from algoliasearch.http.base_config import BaseConfig
+from algoliasearch.http.exceptions import RequestException
+from algoliasearch.http.helpers import (
+    RetryTimeout,
+    create_iterable,
+    create_iterable_sync,
+)
 from algoliasearch.http.request_options import RequestOptions
 from algoliasearch.http.serializer import body_serializer
 from algoliasearch.http.transporter import Transporter
 from algoliasearch.http.transporter_sync import TransporterSync
 from algoliasearch.http.verb import Verb
 from algoliasearch.ingestion.config import IngestionConfig
 from algoliasearch.ingestion.models import (
+    Action,
     ActionType,
     Authentication,
     AuthenticationCreate,
@@ -61,6 +68,7 @@
     OrderKeys,
     PlatformWithNone,
     PushTaskPayload,
+    PushTaskRecords,
     Run,
     RunListResponse,
     RunResponse,
@@ -193,6 +201,75 @@ async def add_user_agent(self, segment: str, version: Optional[str] = None) -> N
         """adds a segment to the default user agent, and update the headers sent with each requests as well"""
         self._transporter.config.add_user_agent(segment, version)
 
+    async def chunked_push(
+        self,
+        index_name: str,
+        objects: List[Dict[str, Any]],
+        action: Action = Action.ADDOBJECT,
+        wait_for_tasks: bool = False,
+        batch_size: int = 1000,
+        reference_index_name: Optional[str] = None,
+        request_options: Optional[Union[dict, RequestOptions]] = None,
+    ) -> List[WatchResponse]:
+        """
+        Helper: Chunks the given `objects` list in subset of 1000 elements max in order to make it fit in `push` requests by leveraging the Transformation pipeline setup in the Push connector (https://www.algolia.com/doc/guides/sending-and-managing-data/send-and-update-your-data/connectors/push/).
+        """
+        records: List[PushTaskRecords] = []
+        responses: List[WatchResponse] = []
+        for i, obj in enumerate(objects):
+            records.append(obj)  # pyright: ignore
+            if len(records) == batch_size or i == len(objects) - 1:
+                responses.append(
+                    await self.push(
+                        index_name=index_name,
+                        push_task_payload={
+                            "action": action,
+                            "records": records,
+                        },
+                        reference_index_name=reference_index_name,
+                        request_options=request_options,
+                    )
+                )
+        if wait_for_tasks:
+            for response in responses:
+
+                async def _func(_: Optional[Event]) -> Event:
+                    if response.event_id is None:
+                        raise ValueError(
+                            "received unexpected response from the push endpoint, eventID must not be undefined"
+                        )
+                    try:
+                        return await self.get_event(
+                            run_id=response.run_id,
+                            event_id=response.event_id,
+                            request_options=request_options,
+                        )
+                    except RequestException as e:
+                        if e.status_code == 404:
+                            return None  # pyright: ignore
+                        raise e
+
+                _retry_count = 0
+
+                def _aggregator(_: Event | None) -> None:
+                    nonlocal _retry_count
+                    _retry_count += 1
+
+                def _validate(_resp: Event | None) -> bool:
+                    return _resp is not None
+
+                timeout = RetryTimeout()
+
+                await create_iterable(
+                    func=_func,
+                    validate=_validate,
+                    aggregator=_aggregator,
+                    timeout=lambda: timeout(_retry_count),
+                    error_validate=lambda _: _retry_count >= 50,
+                    error_message=lambda _: f"The maximum number of retries exceeded. (${_retry_count}/${50})",
+                )
+        return responses
+
     async def create_authentication_with_http_info(
         self,
         authentication_create: Union[AuthenticationCreate, dict[str, Any]],
@@ -5204,6 +5281,75 @@ def add_user_agent(self, segment: str, version: Optional[str] = None) -> None:
         """adds a segment to the default user agent, and update the headers sent with each requests as well"""
         self._transporter.config.add_user_agent(segment, version)
 
+    def chunked_push(
+        self,
+        index_name: str,
+        objects: List[Dict[str, Any]],
+        action: Action = Action.ADDOBJECT,
+        wait_for_tasks: bool = False,
+        batch_size: int = 1000,
+        reference_index_name: Optional[str] = None,
+        request_options: Optional[Union[dict, RequestOptions]] = None,
+    ) -> List[WatchResponse]:
+        """
+        Helper: Chunks the given `objects` list in subset of 1000 elements max in order to make it fit in `push` requests by leveraging the Transformation pipeline setup in the Push connector (https://www.algolia.com/doc/guides/sending-and-managing-data/send-and-update-your-data/connectors/push/).
+        """
+        records: List[PushTaskRecords] = []
+        responses: List[WatchResponse] = []
+        for i, obj in enumerate(objects):
+            records.append(obj)  # pyright: ignore
+            if len(records) == batch_size or i == len(objects) - 1:
+                responses.append(
+                    self.push(
+                        index_name=index_name,
+                        push_task_payload={
+                            "action": action,
+                            "records": records,
+                        },
+                        reference_index_name=reference_index_name,
+                        request_options=request_options,
+                    )
+                )
+        if wait_for_tasks:
+            for response in responses:
+
+                def _func(_: Optional[Event]) -> Event:
+                    if response.event_id is None:
+                        raise ValueError(
+                            "received unexpected response from the push endpoint, eventID must not be undefined"
+                        )
+                    try:
+                        return self.get_event(
+                            run_id=response.run_id,
+                            event_id=response.event_id,
+                            request_options=request_options,
+                        )
+                    except RequestException as e:
+                        if e.status_code == 404:
+                            return None  # pyright: ignore
+                        raise e
+
+                _retry_count = 0
+
+                def _aggregator(_: Event | None) -> None:
+                    nonlocal _retry_count
+                    _retry_count += 1
+
+                def _validate(_resp: Event | None) -> bool:
+                    return _resp is not None
+
+                timeout = RetryTimeout()
+
+                create_iterable_sync(
+                    func=_func,
+                    validate=_validate,
+                    aggregator=_aggregator,
+                    timeout=lambda: timeout(_retry_count),
+                    error_validate=lambda _: _retry_count >= 50,
+                    error_message=lambda _: f"The maximum number of retries exceeded. (${_retry_count}/${50})",
+                )
+        return responses
+
     def create_authentication_with_http_info(
         self,
         authentication_create: Union[AuthenticationCreate, dict[str, Any]],

algoliasearch/search/client.py

@@ -41,6 +41,7 @@
 from algoliasearch.http.verb import Verb
 from algoliasearch.ingestion.client import IngestionClient, IngestionClientSync
 from algoliasearch.ingestion.config import IngestionConfig
+from algoliasearch.ingestion.models import Action as IngestionAction
 from algoliasearch.ingestion.models import WatchResponse
 from algoliasearch.search.config import SearchConfig
 from algoliasearch.search.models import (
@@ -291,7 +292,7 @@ async def wait_for_api_key(
                 "`apiKey` is required when waiting for an `update` operation."
             )
 
-        async def _func(_prev: Optional[GetApiKeyResponse]) -> GetApiKeyResponse:
+        async def _func(_: Optional[GetApiKeyResponse]) -> GetApiKeyResponse:
             try:
                 return await self.get_api_key(key=key, request_options=request_options)
             except RequestException as e:
@@ -431,9 +432,7 @@ async def browse_synonyms(
         page = search_synonyms_params.page or 0
         search_synonyms_params.hits_per_page = hits_per_page
 
-        async def _func(
-            _prev: Optional[SearchSynonymsResponse],
-        ) -> SearchSynonymsResponse:
+        async def _func(_: Optional[SearchSynonymsResponse]) -> SearchSynonymsResponse:
             nonlocal page
             resp = await self.search_synonyms(
                 index_name=index_name,
@@ -534,22 +533,20 @@ async def save_objects_with_transformation(
         wait_for_tasks: bool = False,
         batch_size: int = 1000,
         request_options: Optional[Union[dict, RequestOptions]] = None,
-    ) -> WatchResponse:
+    ) -> List[WatchResponse]:
         """
         Helper: Similar to the `save_objects` method but requires a Push connector (https://www.algolia.com/doc/guides/sending-and-managing-data/send-and-update-your-data/connectors/push/) to be created first, in order to transform records before indexing them to Algolia. The `region` must've been passed to the client's config at instantiation.
         """
         if self._ingestion_transporter is None:
             raise ValueError(
                 "`region` must be provided at client instantiation before calling this method."
             )
-
-        return await self._ingestion_transporter.push(
+        return await self._ingestion_transporter.chunked_push(
             index_name=index_name,
-            push_task_payload={
-                "action": Action.ADDOBJECT,
-                "records": objects,
-            },
-            watch=wait_for_tasks,
+            objects=objects,
+            action=IngestionAction.ADDOBJECT,
+            wait_for_tasks=wait_for_tasks,
+            batch_size=batch_size,
             request_options=request_options,
         )
 
@@ -604,24 +601,22 @@ async def partial_update_objects_with_transformation(
         wait_for_tasks: bool = False,
         batch_size: int = 1000,
         request_options: Optional[Union[dict, RequestOptions]] = None,
-    ) -> WatchResponse:
+    ) -> List[WatchResponse]:
         """
         Helper: Similar to the `partial_update_objects` method but requires a Push connector (https://www.algolia.com/doc/guides/sending-and-managing-data/send-and-update-your-data/connectors/push/) to be created first, in order to transform records before indexing them to Algolia. The `region` must've been passed to the client instantiation method.
         """
         if self._ingestion_transporter is None:
             raise ValueError(
                 "`region` must be provided at client instantiation before calling this method."
             )
-
-        return await self._ingestion_transporter.push(
+        return await self._ingestion_transporter.chunked_push(
             index_name=index_name,
-            push_task_payload={
-                "action": Action.PARTIALUPDATEOBJECT
-                if create_if_not_exists
-                else Action.PARTIALUPDATEOBJECTNOCREATE,
-                "records": objects,
-            },
-            watch=wait_for_tasks,
+            objects=objects,
+            action=IngestionAction.PARTIALUPDATEOBJECT
+            if create_if_not_exists
+            else IngestionAction.PARTIALUPDATEOBJECTNOCREATE,
+            wait_for_tasks=wait_for_tasks,
+            batch_size=batch_size,
             request_options=request_options,
         )
 
@@ -5426,7 +5421,7 @@ def wait_for_api_key(
                 "`apiKey` is required when waiting for an `update` operation."
             )
 
-        def _func(_prev: Optional[GetApiKeyResponse]) -> GetApiKeyResponse:
+        def _func(_: Optional[GetApiKeyResponse]) -> GetApiKeyResponse:
             try:
                 return self.get_api_key(key=key, request_options=request_options)
             except RequestException as e:
@@ -5566,7 +5561,7 @@ def browse_synonyms(
         page = search_synonyms_params.page or 0
         search_synonyms_params.hits_per_page = hits_per_page
 
-        def _func(_prev: Optional[SearchSynonymsResponse]) -> SearchSynonymsResponse:
+        def _func(_: Optional[SearchSynonymsResponse]) -> SearchSynonymsResponse:
             nonlocal page
             resp = self.search_synonyms(
                 index_name=index_name,
@@ -5667,22 +5662,20 @@ def save_objects_with_transformation(
         wait_for_tasks: bool = False,
         batch_size: int = 1000,
         request_options: Optional[Union[dict, RequestOptions]] = None,
-    ) -> WatchResponse:
+    ) -> List[WatchResponse]:
         """
         Helper: Similar to the `save_objects` method but requires a Push connector (https://www.algolia.com/doc/guides/sending-and-managing-data/send-and-update-your-data/connectors/push/) to be created first, in order to transform records before indexing them to Algolia. The `region` must've been passed to the client's config at instantiation.
         """
         if self._ingestion_transporter is None:
             raise ValueError(
                 "`region` must be provided at client instantiation before calling this method."
             )
-
-        return self._ingestion_transporter.push(
+        return self._ingestion_transporter.chunked_push(
             index_name=index_name,
-            push_task_payload={
-                "action": Action.ADDOBJECT,
-                "records": objects,
-            },
-            watch=wait_for_tasks,
+            objects=objects,
+            action=IngestionAction.ADDOBJECT,
+            wait_for_tasks=wait_for_tasks,
+            batch_size=batch_size,
             request_options=request_options,
         )
 
@@ -5737,24 +5730,22 @@ def partial_update_objects_with_transformation(
         wait_for_tasks: bool = False,
         batch_size: int = 1000,
         request_options: Optional[Union[dict, RequestOptions]] = None,
-    ) -> WatchResponse:
+    ) -> List[WatchResponse]:
         """
         Helper: Similar to the `partial_update_objects` method but requires a Push connector (https://www.algolia.com/doc/guides/sending-and-managing-data/send-and-update-your-data/connectors/push/) to be created first, in order to transform records before indexing them to Algolia. The `region` must've been passed to the client instantiation method.
         """
         if self._ingestion_transporter is None:
             raise ValueError(
                 "`region` must be provided at client instantiation before calling this method."
             )
-
-        return self._ingestion_transporter.push(
+        return self._ingestion_transporter.chunked_push(
             index_name=index_name,
-            push_task_payload={
-                "action": Action.PARTIALUPDATEOBJECT
-                if create_if_not_exists
-                else Action.PARTIALUPDATEOBJECTNOCREATE,
-                "records": objects,
-            },
-            watch=wait_for_tasks,
+            objects=objects,
+            action=IngestionAction.PARTIALUPDATEOBJECT
+            if create_if_not_exists
+            else IngestionAction.PARTIALUPDATEOBJECTNOCREATE,
+            wait_for_tasks=wait_for_tasks,
+            batch_size=batch_size,
             request_options=request_options,
         )