This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b358d0a6de Add `retry_args` parameter to `HttpOperator` (#40086)
b358d0a6de is described below
commit b358d0a6de3a7c2fabdabaccf8c5edab9e1d0ecf
Author: Bora Berke Sahin <67373739+borabe...@users.noreply.github.com>
AuthorDate: Thu Jun 6 15:16:30 2024 +0300
Add `retry_args` parameter to `HttpOperator` (#40086)
* Add `retry_args` parameter to `HttpOperator`
* Add unit tests and fix bug in pagination
---
airflow/providers/http/operators/http.py | 18 +++++++-
tests/providers/http/operators/test_http.py | 71 +++++++++++++++++++++++++++++
2 files changed, 87 insertions(+), 2 deletions(-)
diff --git a/airflow/providers/http/operators/http.py
b/airflow/providers/http/operators/http.py
index 3c04a88573..6c086395e5 100644
--- a/airflow/providers/http/operators/http.py
+++ b/airflow/providers/http/operators/http.py
@@ -89,6 +89,8 @@ class HttpOperator(BaseOperator):
:param tcp_keep_alive_interval: The TCP Keep Alive interval parameter
(corresponds to
``socket.TCP_KEEPINTVL``)
:param deferrable: Run operator in the deferrable mode
+ :param retry_args: Arguments which define the retry behaviour.
+ See Tenacity documentation at https://github.com/jd/tenacity
"""
conn_id_field = "http_conn_id"
@@ -120,6 +122,7 @@ class HttpOperator(BaseOperator):
tcp_keep_alive_count: int = 20,
tcp_keep_alive_interval: int = 30,
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ retry_args: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
@@ -139,6 +142,7 @@ class HttpOperator(BaseOperator):
self.tcp_keep_alive_count = tcp_keep_alive_count
self.tcp_keep_alive_interval = tcp_keep_alive_interval
self.deferrable = deferrable
+ self.retry_args = retry_args
@property
def hook(self) -> HttpHook:
@@ -167,7 +171,12 @@ class HttpOperator(BaseOperator):
def execute_sync(self, context: Context) -> Any:
self.log.info("Calling HTTP method")
- response = self.hook.run(self.endpoint, self.data, self.headers,
self.extra_options)
+ if self.retry_args:
+ response = self.hook.run_with_advanced_retry(
+ self.retry_args, self.endpoint, self.data, self.headers,
self.extra_options
+ )
+ else:
+ response = self.hook.run(self.endpoint, self.data, self.headers,
self.extra_options)
response = self.paginate_sync(response=response)
return self.process_response(context=context, response=response)
@@ -180,7 +189,12 @@ class HttpOperator(BaseOperator):
next_page_params = self.pagination_function(response)
if not next_page_params:
break
- response =
self.hook.run(**self._merge_next_page_parameters(next_page_params))
+ if self.retry_args:
+ response = self.hook.run_with_advanced_retry(
+ self.retry_args,
**self._merge_next_page_parameters(next_page_params)
+ )
+ else:
+ response =
self.hook.run(**self._merge_next_page_parameters(next_page_params))
all_responses.append(response)
return all_responses
diff --git a/tests/providers/http/operators/test_http.py
b/tests/providers/http/operators/test_http.py
index dfd82a17ae..3c0d4eb438 100644
--- a/tests/providers/http/operators/test_http.py
+++ b/tests/providers/http/operators/test_http.py
@@ -19,14 +19,18 @@ from __future__ import annotations
import base64
import contextlib
+import json
import pickle
from unittest import mock
+from unittest.mock import call, patch
import pytest
+import tenacity
from requests import Response
from requests.models import RequestEncodingMixin
from airflow.exceptions import AirflowException, TaskDeferred
+from airflow.providers.http.hooks.http import HttpHook
from airflow.providers.http.operators.http import HttpOperator
from airflow.providers.http.triggers.http import HttpTrigger
@@ -228,3 +232,70 @@ class TestHttpOperator:
**create_resume_response_parameters(),
paginated_responses=[make_response_object()]
)
assert result == ['{"value": 5}', '{"value": 5}']
+
+ @patch.object(HttpHook, "run_with_advanced_retry")
+ def test_retry_args(self, mock_run_with_advanced_retry, requests_mock):
+ requests_mock.get("http://www.example.com";, exc=Exception("Example
Exception"))
+ retry_args = dict(
+ wait=tenacity.wait_none(),
+ stop=tenacity.stop_after_attempt(5),
+ retry=tenacity.retry_if_exception_type(Exception),
+ )
+ operator = HttpOperator(
+ task_id="test_HTTP_op",
+ method="GET",
+ endpoint="/",
+ http_conn_id="HTTP_EXAMPLE",
+ retry_args=retry_args,
+ )
+ operator.execute({})
+ mock_run_with_advanced_retry.assert_called_with(retry_args, "/", {},
{}, {})
+ assert mock_run_with_advanced_retry.call_count == 1
+
+ @patch.object(HttpHook, "run_with_advanced_retry")
+ def test_pagination_retry_args(
+ self,
+ mock_run_with_advanced_retry,
+ requests_mock,
+ ):
+ is_second_call: bool = False
+
+ def pagination_function(response: Response) -> dict | None:
+ """Paginated function which returns None at the second call."""
+ nonlocal is_second_call
+ if not is_second_call:
+ is_second_call = True
+ return dict(
+ endpoint=response.json()["endpoint"],
+ )
+ return None
+
+ retry_args = dict(
+ wait=tenacity.wait_none(),
+ stop=tenacity.stop_after_attempt(5),
+ retry=tenacity.retry_if_exception_type(Exception),
+ )
+ operator = HttpOperator(
+ task_id="test_HTTP_op",
+ method="GET",
+ endpoint="/",
+ http_conn_id="HTTP_EXAMPLE",
+ pagination_function=pagination_function,
+ retry_args=retry_args,
+ )
+
+ response = Response()
+ response.status_code = 200
+ response._content = json.dumps({"value": 5, "endpoint":
"/"}).encode("utf-8")
+ response.headers["Content-Type"] = "application/json"
+
+ mock_run_with_advanced_retry.return_value = response
+ operator.execute({})
+ mock_run_with_advanced_retry.assert_has_calls(
+ [
+ call(retry_args, "/", {}, {}, {}),
+ call(retry_args, endpoint="/", data={}, headers={},
extra_options={}),
+ ]
+ )
+
+ assert mock_run_with_advanced_retry.call_count == 2
