Add retry to ESIndex class methods (#2171)

Author: artem-shelkovnikov

connectors/es/index.py

@@ -3,6 +3,8 @@
 # or more contributor license agreements. Licensed under the Elastic License 2.0;
 # you may not use this file except in compliance with the Elastic License 2.0.
 #
+from functools import partial
+
 from elasticsearch import ApiError
 
 from connectors.es import ESClient
@@ -50,10 +52,14 @@ async def fetch_by_id(self, doc_id):
 
     async def fetch_response_by_id(self, doc_id):
         if not self.serverless:
-            await self.client.indices.refresh(index=self.index_name)
+            await self._retrier.execute_with_retry(
+                partial(self.client.indices.refresh, index=self.index_name)
+            )
 
         try:
-            resp = await self.client.get(index=self.index_name, id=doc_id)
+            resp = await self._retrier.execute_with_retry(
+                partial(self.client.get, index=self.index_name, id=doc_id)
+            )
         except ApiError as e:
             logger.critical(f"The server returned {e.status_code}")
             logger.critical(e.body, exc_info=True)
@@ -65,30 +71,41 @@ async def fetch_response_by_id(self, doc_id):
         return resp.body
 
     async def index(self, doc):
-        return await self.client.index(index=self.index_name, document=doc)
+        return await self._retrier.execute_with_retry(
+            partial(self.client.index, index=self.index_name, document=doc)
+        )
 
     async def clean_index(self):
-        return await self.client.delete_by_query(
-            index=self.index_name,
-            body={"query": {"match_all": {}}},
-            ignore_unavailable=True,
-            conflicts="proceed",
+        return await self._retrier.execute_with_retry(
+            partial(
+                self.client.delete_by_query,
+                index=self.index_name,
+                body={"query": {"match_all": {}}},
+                ignore_unavailable=True,
+                conflicts="proceed",
+            )
         )
 
     async def update(self, doc_id, doc, if_seq_no=None, if_primary_term=None):
-        return await self.client.update(
-            index=self.index_name,
-            id=doc_id,
-            doc=doc,
-            if_seq_no=if_seq_no,
-            if_primary_term=if_primary_term,
+        return await self._retrier.execute_with_retry(
+            partial(
+                self.client.update,
+                index=self.index_name,
+                id=doc_id,
+                doc=doc,
+                if_seq_no=if_seq_no,
+                if_primary_term=if_primary_term,
+            )
         )
 
     async def update_by_script(self, doc_id, script):
-        return await self.client.update(
-            index=self.index_name,
-            id=doc_id,
-            script=script,
+        return await self._retrier.execute_with_retry(
+            partial(
+                self.client.update,
+                index=self.index_name,
+                id=doc_id,
+                script=script,
+            )
         )
 
     async def get_all_docs(self, query=None, sort=None, page_size=DEFAULT_PAGE_SIZE):
@@ -103,7 +120,9 @@ async def get_all_docs(self, query=None, sort=None, page_size=DEFAULT_PAGE_SIZE)
             Iterator
         """
         if not self.serverless:
-            await self.client.indices.refresh(index=self.index_name)
+            await self._retrier.execute_with_retry(
+                partial(self.client.indices.refresh, index=self.index_name)
+            )
 
         if query is None:
             query = {"match_all": {}}
@@ -113,14 +132,17 @@ async def get_all_docs(self, query=None, sort=None, page_size=DEFAULT_PAGE_SIZE)
 
         while True:
             try:
-                resp = await self.client.search(
-                    index=self.index_name,
-                    query=query,
-                    sort=sort,
-                    from_=offset,
-                    size=page_size,
-                    expand_wildcards="hidden",
-                    seq_no_primary_term=True,
+                resp = await self._retrier.execute_with_retry(
+                    partial(
+                        self.client.search,
+                        index=self.index_name,
+                        query=query,
+                        sort=sort,
+                        from_=offset,
+                        size=page_size,
+                        expand_wildcards="hidden",
+                        seq_no_primary_term=True,
+                    )
                 )
             except ApiError as e:
                 logger.error(

tests/es/test_index.py

@@ -118,7 +118,7 @@ async def test_fetch_response_by_id_not_found(mock_responses):
 
 
 @pytest.mark.asyncio
-async def test_fetch_response_by_id_api_error(mock_responses):
+async def test_fetch_response_by_id_api_error(mock_responses, patch_sleep):
     doc_id = "1"
     index = FakeIndex(index_name, config)
     mock_responses.post(
@@ -129,6 +129,7 @@ async def test_fetch_response_by_id_api_error(mock_responses):
         f"http://nowhere.com:9200/{index_name}/_doc/{doc_id}",
         headers=headers,
         status=500,
+        repeat=True,
     )
 
     with pytest.raises(ApiError):
@@ -201,7 +202,7 @@ async def test_update_by_script():
 
 
 @pytest.mark.asyncio
-async def test_get_all_docs_with_error(mock_responses, patch_logger):
+async def test_get_all_docs_with_error(mock_responses, patch_logger, patch_sleep):
     index = FakeIndex(index_name, config)
     mock_responses.post(
         f"http://nowhere.com:9200/{index_name}/_refresh", headers=headers, status=200
@@ -211,6 +212,7 @@ async def test_get_all_docs_with_error(mock_responses, patch_logger):
         f"http://nowhere.com:9200/{index_name}/_search?expand_wildcards=hidden",
         headers=headers,
         status=500,
+        repeat=True,
     )
 
     with pytest.raises(ApiError) as e: