Add pagination to SimpleHttpOperator

airflow/providers/http/operators/http.py

@@ -21,14 +21,16 @@
 import pickle
 from typing import TYPE_CHECKING, Any, Callable, Sequence
 
+from requests import Response
+
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.providers.http.hooks.http import HttpHook
 from airflow.providers.http.triggers.http import HttpTrigger
+from airflow.utils.helpers import merge_dicts
 
 if TYPE_CHECKING:
-    from requests import Response
     from requests.auth import AuthBase
 
     from airflow.utils.context import Context
@@ -49,14 +51,25 @@ class SimpleHttpOperator(BaseOperator):
     :param data: The data to pass. POST-data in POST/PUT and params
         in the URL for a GET request. (templated)
     :param headers: The HTTP headers to be added to the GET request
+    :param pagination_function: A callable that generate the parameters used to call the API again.
+        Typically used when the API is paginated and returns for e.g a cursor, a 'next page id', or
+        a 'next page URL'. When provided, the Operator will call the API repeatedly until this function
+        returns None. Also, the result of the Operator will be by default a list of Response.text
+        objects instead of a single response object. This function should return a dict of parameters
+        (`endpoint`, `data`, `headers`, `extra_options`), which will merge and override the one used in
+        the initial or previous API call.
     :param response_check: A check against the 'requests' response object.
         The callable takes the response object as the first positional argument
         and optionally any number of keyword arguments available in the context dictionary.
-        It should return True for 'pass' and False otherwise.
+        It should return True for 'pass' and False otherwise. If a pagination_function
+        is provided, this function will receive a list  of response object instead of a
+        single response object.
     :param response_filter: A function allowing you to manipulate the response
         text. e.g response_filter=lambda response: json.loads(response.text).
         The callable takes the response object as the first positional argument
         and optionally any number of keyword arguments available in the context dictionary.
+        If a pagination_function is provided, this function will receive a list of response
+        object instead of a single response object.
     :param extra_options: Extra options for the 'requests' library, see the
         'requests' documentation (options to modify timeout, ssl, etc.)
     :param log_response: Log the response (default: False)
@@ -85,6 +98,7 @@ def __init__(
         method: str = "POST",
         data: Any = None,
         headers: dict[str, str] | None = None,
+        pagination_function: Callable[..., Any] | None = None,
         response_check: Callable[..., bool] | None = None,
         response_filter: Callable[..., Any] | None = None,
         extra_options: dict[str, Any] | None = None,
@@ -104,6 +118,7 @@ def __init__(
         self.endpoint = endpoint
         self.headers = headers or {}
         self.data = data or {}
+        self.pagination_function = pagination_function
         self.response_check = response_check
         self.response_filter = response_filter
         self.extra_options = extra_options or {}
@@ -141,33 +156,97 @@ def execute(self, context: Context) -> Any:
             )
 
             self.log.info("Calling HTTP method")
-
             response = http.run(self.endpoint, self.data, self.headers, self.extra_options)
+
+            if self.pagination_function:
+                all_responses: list[Response] = [response]
+                while True:
+                    next_page_params = self.pagination_function(response)
+                    if not next_page_params:
+                        break
+                    response = http.run(**self._merge_next_page_parameters(next_page_params))
+                    all_responses.append(response)
+                response = all_responses
+
             return self.process_response(context=context, response=response)
 
-    def process_response(self, context: Context, response: Response) -> str:
+    def process_response(self, context: Context, response: Response | list[Response]) -> str | list[str]:
         """Process the response."""
         from airflow.utils.operator_helpers import determine_kwargs
 
+        make_default_response: Callable = self._default_response_maker(response=response)
+
         if self.log_response:
-            self.log.info(response.text)
+            self.log.info(make_default_response())
         if self.response_check:
             kwargs = determine_kwargs(self.response_check, [response], context)
             if not self.response_check(response, **kwargs):
                 raise AirflowException("Response check returned False.")
         if self.response_filter:
             kwargs = determine_kwargs(self.response_filter, [response], context)
             return self.response_filter(response, **kwargs)
-        return response.text
+        return make_default_response()
+
+    @staticmethod
+    def _default_response_maker(response: Response | list[Response]) -> Callable:
+        """Create a default response maker function based on the type of response.
 
-    def execute_complete(self, context: Context, event: dict):
+        :param response: The response object or list of response objects.
+        :return: A function that returns response text(s).
         """
-        Callback for when the trigger fires - returns immediately.
+        if isinstance(response, Response):
+            response_object = response  # Makes mypy happy
+            return lambda: response_object.text
+
+        response_list: list[Response] = response  # Makes mypy happy
+        return lambda: [entry.text for entry in response_list]
 
+    def execute_complete(
+        self, context: Context, event: dict, paginated_responses: None | list[Response] = None
+    ):
+        """Callback for when the trigger fires.
+
+        When no pagination, this method returns immediately. Otherwise, it creates a new deferrable.
         Relies on trigger to throw an exception, otherwise it assumes execution was successful.
         """
         if event["status"] == "success":
             response = pickle.loads(base64.standard_b64decode(event["response"]))
-            return self.process_response(context=context, response=response)
+
+            if self.pagination_function:
+                paginated_responses = paginated_responses or []
+                paginated_responses.append(response)
+
+                next_page_params = self.pagination_function(response)
+                if not next_page_params:
+                    return self.process_response(context=context, response=paginated_responses)
+                self.defer(
+                    trigger=HttpTrigger(
+                        http_conn_id=self.http_conn_id,
+                        auth_type=self.auth_type,
+                        method=self.method,
+                        **self._merge_next_page_parameters(next_page_params),
+                    ),
+                    method_name="execute_complete",
+                    kwargs={"paginated_responses": paginated_responses},
+                )
+            else:
+                return self.process_response(context=context, response=response)
         else:
             raise AirflowException(f"Unexpected error in the operation: {event['message']}")
+
+    def _merge_next_page_parameters(self, next_page_params: dict) -> dict:
+        """Merge initial request parameters with next page parameters.
+
+        Merge initial requests parameters with the ones for the next page, generated by
+        the pagination function. Items in the 'next_page_params' overrides those defined
+        in the previous request.
+
+        :param next_page_params: A dictionary containing the parameters for the next page.
+        :return: A dictionary containing the merged parameters.
+        """
+        return dict(
+            endpoint=next_page_params.get("endpoint") or self.endpoint,
+            data=merge_dicts(self.data, next_page_params.get("data", {})),
+            headers=merge_dicts(self.headers, next_page_params.get("headers", {})),
+            extra_options=merge_dicts(self.extra_options, next_page_params.get("extra_options", {})),
+        )

docs/apache-airflow-providers-http/operators.rst

@@ -96,6 +96,18 @@ response body.
     :start-after: [START howto_operator_http_task_get_op_response_filter]
     :end-before: [END howto_operator_http_task_get_op_response_filter]
 
+You can also do a series of call on a paginated API. Let's assume your API returns a JSON body containing a cursor:
+You can write a ``pagination_function`` that will receive the raw ``request.Response`` object of your request, and
+generate new request parameters (as ``dict``) based on this cursor. The SimpleHttpOperator will repeat calls to the
+API until the function stop returning anything.
+
+In this case, by default, the result of the SimpleHttpOperator will be a list of text.
+
+.. exampleinclude:: /../../tests/system/providers/http/example_http.py
+    :language: python
+    :start-after: [START howto_operator_http_pagination_function]
+    :end-before: [END howto_operator_http_pagination_function]
+
 In the third example we are performing a ``PUT`` operation to put / set data according to the data that is being
 provided to the request.
 

tests/providers/http/operators/test_http.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import base64
+import contextlib
 import pickle
 from unittest import mock
 
@@ -86,6 +87,37 @@ def test_filters_response(self, requests_mock):
         result = operator.execute({})
         assert result == {"value": 5}
 
+    def test_paginated_responses(self, requests_mock):
+        """
+        Test that the SimpleHttpOperator calls repetitively the API when a
+        pagination_function is provided, and as long as this function returns
+        a dictionary that override previous' call parameters.
+        """
+        has_returned: bool = False
+
+        def pagination_function(response: Response) -> dict | None:
+            """Paginated function which returns None at the second call."""
+            nonlocal has_returned
+            if not has_returned:
+                has_returned = True
+                return dict(
+                    endpoint="/",
+                    data={"cursor": "example"},
+                    headers={},
+                    extra_options={},
+                )
+
+        requests_mock.get("http://www.example.com", json={"value": 5})
+        operator = SimpleHttpOperator(
+            task_id="test_HTTP_op",
+            method="GET",
+            endpoint="/",
+            http_conn_id="HTTP_EXAMPLE",
+            pagination_function=pagination_function,
+        )
+        result = operator.execute({})
+        assert result == ['{"value": 5}', '{"value": 5}']
+
     def test_async_defer_successfully(self, requests_mock):
         operator = SimpleHttpOperator(
             task_id="test_HTTP_op",
@@ -110,3 +142,50 @@ def test_async_execute_successfully(self, requests_mock):
             },
         )
         assert result == "content"
+
+    def test_async_paginated_responses(self, requests_mock):
+        """
+        Test that the SimpleHttpOperator calls asynchronously and repetitively
+        the API when a pagination_function is provided, and as long as this function
+        returns a dictionary that override previous' call parameters.
+        """
+
+        def make_response_object() -> Response:
+            response = Response()
+            response._content = b'{"value": 5}'
+            return response
+
+        def create_resume_response_parameters() -> dict:
+            response = make_response_object()
+            return dict(
+                context={},
+                event={
+                    "status": "success",
+                    "response": base64.standard_b64encode(pickle.dumps(response)).decode("ascii"),
+                },
+            )
+
+        has_returned: bool = False
+
+        def pagination_function(response: Response) -> dict | None:
+            """Paginated function which returns None at the second call."""
+            nonlocal has_returned
+            if not has_returned:
+                has_returned = True
+                return dict(endpoint="/")
+
+        operator = SimpleHttpOperator(
+            task_id="test_HTTP_op",
+            pagination_function=pagination_function,
+            deferrable=True,
+        )
+
+        # Do two calls: On the first one, the pagination_function creates a new
+        # deferrable trigger. On the second one, the pagination_function returns
+        # None, which ends the execution of the Operator
+        with contextlib.suppress(TaskDeferred):
+            operator.execute_complete(**create_resume_response_parameters())
+            result = operator.execute_complete(
+                **create_resume_response_parameters(), paginated_responses=[make_response_object()]
+            )
+            assert result == ['{"value": 5}', '{"value": 5}']

tests/system/providers/http/example_http.py

@@ -110,8 +110,32 @@
     dag=dag,
 )
 # [END howto_operator_http_http_sensor_check]
+# [START howto_operator_http_pagination_function]
+
+
+def get_next_page_cursor(response) -> dict | None:
+    """
+    Take the raw `request.Response` object, and check for a cursor.
+    If a cursor exists, this function creates and return parameters to call
+    the next page of result.
+    """
+    next_cursor = response.json().get("cursor")
+    if next_cursor:
+        return dict(data={"cursor": next_cursor})
+
+
+task_get_paginated = SimpleHttpOperator(
+    task_id="get_paginated",
+    method="GET",
+    endpoint="get",
+    data={"cursor": ""},
+    pagination_function=get_next_page_cursor,
+    dag=dag,
+)
+# [END howto_operator_http_pagination_function]
 task_http_sensor_check >> task_post_op >> task_get_op >> task_get_op_response_filter
 task_get_op_response_filter >> task_put_op >> task_del_op >> task_post_op_formenc
+task_post_op_formenc >> task_get_paginated
 
 from tests.system.utils import get_test_run  # noqa: E402