main commit

venv/lib/python3.12/site-packages/celery/backends/elasticsearch.py

@@ -1,5 +1,5 @@
 """Elasticsearch result store backend."""
-from datetime import datetime, timezone
+from datetime import datetime
 
 from kombu.utils.encoding import bytes_to_str
 from kombu.utils.url import _parse_url
@@ -14,11 +14,6 @@ try:
 except ImportError:
     elasticsearch = None
 
-try:
-    import elastic_transport
-except ImportError:
-    elastic_transport = None
-
 __all__ = ('ElasticsearchBackend',)
 
 E_LIB_MISSING = """\
@@ -36,7 +31,7 @@ class ElasticsearchBackend(KeyValueStoreBackend):
     """
 
     index = 'celery'
-    doc_type = None
+    doc_type = 'backend'
     scheme = 'http'
     host = 'localhost'
     port = 9200
@@ -88,17 +83,17 @@ class ElasticsearchBackend(KeyValueStoreBackend):
         self._server = None
 
     def exception_safe_to_retry(self, exc):
-        if isinstance(exc, elasticsearch.exceptions.ApiError):
+        if isinstance(exc, (elasticsearch.exceptions.TransportError)):
             # 401: Unauthorized
             # 409: Conflict
+            # 429: Too Many Requests
             # 500: Internal Server Error
             # 502: Bad Gateway
+            # 503: Service Unavailable
             # 504: Gateway Timeout
             # N/A: Low level exception (i.e. socket exception)
-            if exc.status_code in {401, 409, 500, 502, 504, 'N/A'}:
+            if exc.status_code in {401, 409, 429, 500, 502, 503, 504, 'N/A'}:
                 return True
-        if isinstance(exc, elasticsearch.exceptions.TransportError):
-            return True
         return False
 
     def get(self, key):
@@ -113,23 +108,17 @@ class ElasticsearchBackend(KeyValueStoreBackend):
             pass
 
     def _get(self, key):
-        if self.doc_type:
-            return self.server.get(
-                index=self.index,
-                id=key,
-                doc_type=self.doc_type,
-            )
-        else:
-            return self.server.get(
-                index=self.index,
-                id=key,
-            )
+        return self.server.get(
+            index=self.index,
+            doc_type=self.doc_type,
+            id=key,
+        )
 
     def _set_with_state(self, key, value, state):
         body = {
             'result': value,
             '@timestamp': '{}Z'.format(
-                datetime.now(timezone.utc).isoformat()[:-9]
+                datetime.utcnow().isoformat()[:-3]
             ),
         }
         try:
@@ -146,23 +135,14 @@ class ElasticsearchBackend(KeyValueStoreBackend):
 
     def _index(self, id, body, **kwargs):
         body = {bytes_to_str(k): v for k, v in body.items()}
-        if self.doc_type:
-            return self.server.index(
-                id=bytes_to_str(id),
-                index=self.index,
-                doc_type=self.doc_type,
-                body=body,
-                params={'op_type': 'create'},
-                **kwargs
-            )
-        else:
-            return self.server.index(
-                id=bytes_to_str(id),
-                index=self.index,
-                body=body,
-                params={'op_type': 'create'},
-                **kwargs
-            )
+        return self.server.index(
+            id=bytes_to_str(id),
+            index=self.index,
+            doc_type=self.doc_type,
+            body=body,
+            params={'op_type': 'create'},
+            **kwargs
+        )
 
     def _update(self, id, body, state, **kwargs):
         """Update state in a conflict free manner.
@@ -202,32 +182,19 @@ class ElasticsearchBackend(KeyValueStoreBackend):
         prim_term = res_get.get('_primary_term', 1)
 
         # try to update document with current seq_no and primary_term
-        if self.doc_type:
-            res = self.server.update(
-                id=bytes_to_str(id),
-                index=self.index,
-                doc_type=self.doc_type,
-                body={'doc': body},
-                params={'if_primary_term': prim_term, 'if_seq_no': seq_no},
-                **kwargs
-            )
-        else:
-            res = self.server.update(
-                id=bytes_to_str(id),
-                index=self.index,
-                body={'doc': body},
-                params={'if_primary_term': prim_term, 'if_seq_no': seq_no},
-                **kwargs
-            )
+        res = self.server.update(
+            id=bytes_to_str(id),
+            index=self.index,
+            doc_type=self.doc_type,
+            body={'doc': body},
+            params={'if_primary_term': prim_term, 'if_seq_no': seq_no},
+            **kwargs
+        )
         # result is elastic search update query result
         # noop = query did not update any document
         # updated = at least one document got updated
         if res['result'] == 'noop':
-            raise elasticsearch.exceptions.ConflictError(
-                "conflicting update occurred concurrently",
-                elastic_transport.ApiResponseMeta(409, "HTTP/1.1",
-                                                  elastic_transport.HttpHeaders(), 0, elastic_transport.NodeConfig(
-                                                      self.scheme, self.host, self.port)), None)
+            raise elasticsearch.exceptions.ConflictError(409, 'conflicting update occurred concurrently', {})
         return res
 
     def encode(self, data):
@@ -258,10 +225,7 @@ class ElasticsearchBackend(KeyValueStoreBackend):
         return [self.get(key) for key in keys]
 
     def delete(self, key):
-        if self.doc_type:
-            self.server.delete(index=self.index, id=key, doc_type=self.doc_type)
-        else:
-            self.server.delete(index=self.index, id=key)
+        self.server.delete(index=self.index, doc_type=self.doc_type, id=key)
 
     def _get_server(self):
         """Connect to the Elasticsearch server."""
@@ -269,10 +233,11 @@ class ElasticsearchBackend(KeyValueStoreBackend):
         if self.username and self.password:
             http_auth = (self.username, self.password)
         return elasticsearch.Elasticsearch(
-            f'{self.scheme}://{self.host}:{self.port}',
+            f'{self.host}:{self.port}',
             retry_on_timeout=self.es_retry_on_timeout,
             max_retries=self.es_max_retries,
             timeout=self.es_timeout,
+            scheme=self.scheme,
             http_auth=http_auth,
         )