diff --git a/CHANGELOG.md b/CHANGELOG.md index 29b7560..7aac3bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ## [Version 1.4.3](https://github.com/dataiku/dss-plugin-api-connect/releases/tag/v1.4.3) - Bugfix - 2026-07-27 - Fix templating for multiform body +- Adding a configurable retry for several HTTP errors ## [Version 1.4.2](https://github.com/dataiku/dss-plugin-api-connect/releases/tag/v1.4.2) - Bugfix - 2026-07-22 diff --git a/custom-recipes/api-connect/recipe.json b/custom-recipes/api-connect/recipe.json index 2a1a757..f943907 100644 --- a/custom-recipes/api-connect/recipe.json +++ b/custom-recipes/api-connect/recipe.json @@ -408,6 +408,77 @@ "description": "-1 for no limit", "type": "INT", "defaultValue": -1 + }, + { + "name": "http_errors_retry_strategy", + "label": "Retry on error logic", + "description": "", + "type": "SELECT", + "defaultValue": null, + "selectChoices":[ + {"value": null, "label": "No retry"}, + {"value": "linear", "label": "Linear backoff"}, + {"value": "exponential", "label": "Exponential backoff"} + ] + }, + { + "name": "http_errors_to_retry", + "label": "Errors to retry", + "description": "Click to select errors that can trigger a retry", + "type": "MULTISELECT", + "defaultValue": null, + "selectChoices":[ + {"value": "408", "label": "408 Request Timeout"}, + {"value": "429", "label": "429 Too many requests"}, + {"value": "502", "label": "502 Bad Gateway"}, + {"value": "503", "label": "503 Service Unavailable"}, + {"value": "504", "label": "504 Gateway Time out"} + ], + "visibilityCondition": "(['exponential', 'linear'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_retry_scope", + "label": "Retry scope", + "description": "Apply the retry budget to the entire input dataset or independently to each input row", + "type": "SELECT", + "defaultValue": "dataset", + "selectChoices":[ + {"value": "dataset", "label": "Per dataset"}, + {"value": "row", "label": "Per row"} + ], + "visibilityCondition": "(['exponential', 'linear'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_initial_delay", + "label": "Initial delay", + "description": "in seconds", + "type": "INT", + "defaultValue": 1, + "visibilityCondition": "(['exponential'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_maximum_delay", + "label": "Maximum delay", + "description": "in seconds", + "type": "INT", + "defaultValue": 120, + "visibilityCondition": "(['exponential'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_delay", + "label": "Delay", + "description": "in seconds", + "type": "INT", + "defaultValue": 1, + "visibilityCondition": "(['linear'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_maximum_retries", + "label": "Maximum number of retries", + "description": "Number of times to retry a request after an error", + "type": "INT", + "defaultValue": 5, + "visibilityCondition": "(['linear'].includes(model.http_errors_retry_strategy))" } ], "resourceKeys": [] diff --git a/custom-recipes/api-connect/recipe.py b/custom-recipes/api-connect/recipe.py index 54fd9ad..f119798 100644 --- a/custom-recipes/api-connect/recipe.py +++ b/custom-recipes/api-connect/recipe.py @@ -3,9 +3,10 @@ from dataiku.customrecipe import get_input_names_for_role, get_recipe_config, get_output_names_for_role import pandas as pd from safe_logger import SafeLogger -from dku_utils import get_dku_key_values, get_endpoint_parameters, get_secure_credentials, get_user_secrets +from dku_utils import get_dku_key_values, get_endpoint_parameters, get_secure_credentials, get_user_secrets, get_retry_handler_parameters_from_config from rest_api_recipe_session import RestApiRecipeSession from dku_constants import DKUConstants +from retry_handler import RetryHandler logger = SafeLogger("api-connect plugin", forbidden_keys=DKUConstants.FORBIDDEN_KEYS) @@ -49,10 +50,18 @@ def get_partitioning_keys(id_list, dku_flow_variables): custom_key_values.update(user_secrets) display_metadata = config.get("display_metadata", False) maximum_number_rows = config.get("maximum_number_rows", -1) +retry_scope = config.get("http_errors_retry_scope", "dataset") input_parameters_dataset = dataiku.Dataset(input_A_names[0]) partitioning_keys = get_partitioning_keys(input_parameters_dataset, dku_flow_variables) custom_key_values.update(partitioning_keys) input_parameters_dataframe = input_parameters_dataset.get_dataframe(infer_with_pandas=False) +backoff_type, initial_delay, maximum_number_of_retries, maximum_duration_of_retry, status_codes_to_retry = get_retry_handler_parameters_from_config(config) +retry_handler = None +if backoff_type: + retry_handler = RetryHandler( + backoff_type=backoff_type, initial_delay=initial_delay, maximum_number_of_retries=maximum_number_of_retries, + maximum_duration_of_retry=maximum_duration_of_retry, status_codes_to_retry=status_codes_to_retry + ) recipe_session = RestApiRecipeSession( custom_key_values, @@ -64,7 +73,9 @@ def get_partitioning_keys(id_list, dku_flow_variables): parameter_renamings, display_metadata, maximum_number_rows=maximum_number_rows, - behaviour_when_error=behaviour_when_error + behaviour_when_error=behaviour_when_error, + retry_handler=retry_handler, + retry_scope=retry_scope ) results = recipe_session.process_dataframe(input_parameters_dataframe, is_raw_output) diff --git a/python-connectors/api-connect_dataset/connector.json b/python-connectors/api-connect_dataset/connector.json index db18a68..ce965ec 100644 --- a/python-connectors/api-connect_dataset/connector.json +++ b/python-connectors/api-connect_dataset/connector.json @@ -356,6 +356,65 @@ "description": "-1 for no limit", "type": "INT", "defaultValue": -1 + }, + { + "name": "http_errors_retry_strategy", + "label": "Retry on error logic", + "description": "", + "type": "SELECT", + "defaultValue": null, + "selectChoices":[ + {"value": null, "label": "No retry"}, + {"value": "linear", "label": "Linear backoff"}, + {"value": "exponential", "label": "Exponential backoff"} + ] + }, + { + "name": "http_errors_to_retry", + "label": "Errors to retry", + "description": "Click to select errors that can trigger a retry", + "type": "MULTISELECT", + "defaultValue": null, + "selectChoices":[ + {"value": "408", "label": "408 Request Timeout"}, + {"value": "429", "label": "429 Too many requests"}, + {"value": "502", "label": "502 Bad Gateway"}, + {"value": "503", "label": "503 Service Unavailable"}, + {"value": "504", "label": "504 Gateway Time out"} + ], + "visibilityCondition": "(['exponential', 'linear'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_initial_delay", + "label": "Initial delay", + "description": "in seconds", + "type": "INT", + "defaultValue": 1, + "visibilityCondition": "(['exponential'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_maximum_delay", + "label": "Maximum delay", + "description": "in seconds", + "type": "INT", + "defaultValue": 120, + "visibilityCondition": "(['exponential'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_delay", + "label": "Delay", + "description": "in seconds", + "type": "INT", + "defaultValue": 1, + "visibilityCondition": "(['linear'].includes(model.http_errors_retry_strategy))" + }, + { + "name": "http_errors_maximum_retries", + "label": "Maximum number of retries", + "description": "Number of times to retry a request after an error", + "type": "INT", + "defaultValue": 5, + "visibilityCondition": "(['linear'].includes(model.http_errors_retry_strategy))" } ] } diff --git a/python-connectors/api-connect_dataset/connector.py b/python-connectors/api-connect_dataset/connector.py index baff370..3e88c26 100644 --- a/python-connectors/api-connect_dataset/connector.py +++ b/python-connectors/api-connect_dataset/connector.py @@ -5,10 +5,11 @@ from dku_utils import ( get_dku_key_values, get_endpoint_parameters, parse_keys_for_json, get_value_from_path, get_secure_credentials, - decode_csv_data, decode_bytes, get_user_secrets + decode_csv_data, decode_bytes, get_user_secrets, get_retry_handler_parameters_from_config ) from dku_constants import DKUConstants import json +from retry_handler import RetryHandler logger = SafeLogger("api-connect plugin", forbidden_keys=DKUConstants.FORBIDDEN_KEYS) @@ -26,7 +27,14 @@ def __init__(self, config, plugin_config): custom_key_values = get_dku_key_values(config.get("custom_key_values", {})) user_secrets = get_user_secrets(config) custom_key_values.update(user_secrets) - self.client = RestAPIClient(credential, secure_credentials, endpoint_parameters, custom_key_values) + backoff_type, initial_delay, maximum_number_of_retries, maximum_duration_of_retry, status_codes_to_retry = get_retry_handler_parameters_from_config(config) + retry_handler = None + if backoff_type: + retry_handler = RetryHandler( + backoff_type=backoff_type, initial_delay=initial_delay, maximum_number_of_retries=maximum_number_of_retries, + maximum_duration_of_retry=maximum_duration_of_retry, status_codes_to_retry=status_codes_to_retry + ) + self.client = RestAPIClient(credential, secure_credentials, endpoint_parameters, custom_key_values, retry_handler=retry_handler) extraction_key = endpoint_parameters.get("extraction_key", None) self.extraction_key = extraction_key or '' self.extraction_path = self.extraction_key.split('.') diff --git a/python-lib/dku_utils.py b/python-lib/dku_utils.py index 692300a..498238b 100644 --- a/python-lib/dku_utils.py +++ b/python-lib/dku_utils.py @@ -319,3 +319,21 @@ def join_url(base_url, segment): segment = segment.lstrip("/") segments.append(segment) return "/".join(segments) + + +def get_retry_handler_parameters_from_config(config): + backoff_type = initial_delay = maximum_number_of_retries = maximum_duration_of_retry = status_codes_to_retry = None + http_errors_retry_strategy = config.get("http_errors_retry_strategy", None) + if not http_errors_retry_strategy: + return backoff_type, initial_delay, maximum_number_of_retries, maximum_duration_of_retry, status_codes_to_retry + if http_errors_retry_strategy in ["linear", "exponential"]: + backoff_type = http_errors_retry_strategy + if backoff_type == "linear": + initial_delay = config.get("http_errors_delay") + maximum_number_of_retries = config.get("http_errors_maximum_retries", None) + if backoff_type == "exponential": + initial_delay = config.get("http_errors_initial_delay") + maximum_duration_of_retry = config.get("http_errors_maximum_delay", None) + if backoff_type: + status_codes_to_retry = config.get("http_errors_to_retry", []) + return backoff_type, initial_delay, maximum_number_of_retries, maximum_duration_of_retry, status_codes_to_retry diff --git a/python-lib/rest_api_client.py b/python-lib/rest_api_client.py index d57e5d4..0d86d91 100644 --- a/python-lib/rest_api_client.py +++ b/python-lib/rest_api_client.py @@ -8,6 +8,7 @@ from dku_utils import get_dku_key_values, get_dku_duplicated_key_values, template_dict, format_template, is_reponse_xml, xml_to_json from dku_constants import DKUConstants from rest_api_auth import get_auth +from retry_handler import DefaultRetryHandler logger = SafeLogger("api-connect plugin", forbidden_keys=DKUConstants.FORBIDDEN_KEYS) @@ -19,7 +20,7 @@ class RestAPIClientError(ValueError): class RestAPIClient(object): - def __init__(self, credential, secure_credentials, endpoint, custom_key_values={}, session=None, behaviour_when_error=None): + def __init__(self, credential, secure_credentials, endpoint, custom_key_values={}, session=None, behaviour_when_error=None, retry_handler=None): logger.info("Initialising RestAPIClient, credential={}, secure_credentials={}, endpoint={}".format( logger.filter_secrets(credential), logger.filter_secrets(secure_credentials), @@ -134,6 +135,7 @@ def __init__(self, credential, secure_credentials, endpoint, custom_key_values={ self.secure_domain = "https://{}".format(self.secure_domain) else: self.session.auth = get_auth(credential) + self.retry_handler = retry_handler or DefaultRetryHandler() def get(self, url, can_raise_exeption=True, **kwargs): json_response = self.request("GET", url, can_raise_exeption=can_raise_exeption, **kwargs) @@ -142,12 +144,10 @@ def get(self, url, can_raise_exeption=True, **kwargs): def request(self, method, url, can_raise_exeption=True, **kwargs): logger.info(u"Accessing endpoint {} with params={}".format(url, kwargs.get("params"))) self.assert_secure_domain(url) - self.enforce_throttling() kwargs = template_dict(kwargs, **self.presets_variables) if self.loop_detector.is_stuck_in_loop(url, kwargs.get("params", {}), kwargs.get("headers", {})): raise RestAPIClientError("The api-connect plugin is stuck in a loop. Please check the pagination parameters.") request_start_time = time.time() - self.time_last_request = request_start_time error_message = None status_code = None response_headers = None @@ -216,9 +216,17 @@ def request_with_cert(self, method, url, **kwargs): ) tmp_key.seek(0) kwargs["cert"] = (tmp_certificate.name, tmp_key.name) - response = self.session.request(method, url, **kwargs) + response = self.request_with_errors_retry(method, url, **kwargs) return response - return self.session.request(method, url, **kwargs) + return self.request_with_errors_retry(method, url, **kwargs) + + def request_with_errors_retry(self, method, url, **kwargs): + response = None + while self.retry_handler.should_retry(response): + self.enforce_throttling() + self.time_last_request = time.time() + response = self.session.request(method, url, **kwargs) + return response def paginated_api_call(self, can_raise_exeption=True): if self.pagination.params_must_be_blanked: @@ -256,7 +264,7 @@ def start_paging(self): self.pagination.reset_paging(counting_key=self.extraction_key, url=self.endpoint_url) def enforce_throttling(self): - if self.time_between_requests and self.time_last_request: + if self.time_between_requests and self.time_last_request is not None: current_time = time.time() time_since_last_resquests = current_time - self.time_last_request if time_since_last_resquests < self.time_between_requests: diff --git a/python-lib/rest_api_recipe_session.py b/python-lib/rest_api_recipe_session.py index f8ef583..e9e9d05 100644 --- a/python-lib/rest_api_recipe_session.py +++ b/python-lib/rest_api_recipe_session.py @@ -15,7 +15,7 @@ class RestApiRecipeSession: def __init__(self, custom_key_values, credential_parameters, secure_credentials, endpoint_parameters, extraction_key, parameter_columns, parameter_renamings, display_metadata=False, - maximum_number_rows=-1, behaviour_when_error=None): + maximum_number_rows=-1, behaviour_when_error=None, retry_handler=None, retry_scope="dataset"): self.custom_key_values = custom_key_values self.credential_parameters = credential_parameters self.secure_credentials = secure_credentials @@ -30,6 +30,8 @@ def __init__(self, custom_key_values, credential_parameters, secure_credentials, self.behaviour_when_error = behaviour_when_error or "add-error-column" self.can_raise = self.behaviour_when_error == "raise" self.csv_configuration = endpoint_parameters + self.retry_handler = retry_handler + self.retry_scope = retry_scope @staticmethod def get_column_to_parameter_dict(parameter_columns, parameter_renamings): @@ -46,6 +48,9 @@ def process_dataframe(self, input_parameters_dataframe, is_raw_output): time_last_request = None session = requests.Session() for index, input_parameters_row in input_parameters_dataframe.iterrows(): + retry_handler = self.retry_handler + if self.retry_scope == "row" and retry_handler: + retry_handler = retry_handler.recreate() rows_count = 0 self.initial_parameter_columns = {} for column_name in self.column_to_parameter_dict: @@ -68,7 +73,8 @@ def process_dataframe(self, input_parameters_dataframe, is_raw_output): updated_endpoint_parameters, custom_key_values=self.custom_key_values, session=session, - behaviour_when_error=self.behaviour_when_error + behaviour_when_error=self.behaviour_when_error, + retry_handler=retry_handler ) self.client.time_last_request = time_last_request while self.client.has_more_data(): diff --git a/python-lib/retry_handler.py b/python-lib/retry_handler.py new file mode 100644 index 0000000..c38585f --- /dev/null +++ b/python-lib/retry_handler.py @@ -0,0 +1,107 @@ +import requests +import time +from safe_logger import SafeLogger + + +logger = SafeLogger("api-connect plugin retry handler") + + +class DefaultRetryHandler(): + def __init__(self): + pass + + def should_retry(self, response): + if response is None: + return True + return False + + +class RetryHandler(): + def __init__(self, backoff_type=None, initial_delay=None, maximum_number_of_retries=None, + maximum_duration_of_retry=None, status_codes_to_retry=None): + self.backoff_type = None + if backoff_type in ["linear", "exponential"]: + self.backoff_type = backoff_type + self.initial_delay = 0 + if isinstance(initial_delay, int): + self.initial_delay = initial_delay + self.maximum_number_of_retries = None + if isinstance(maximum_number_of_retries, int): + self.maximum_number_of_retries = maximum_number_of_retries + self.maximum_duration_of_retry = None + if isinstance(maximum_duration_of_retry, int): + self.maximum_duration_of_retry = maximum_duration_of_retry + self.next_delay = None + self.status_codes_to_retry = [] + if isinstance(status_codes_to_retry, list): + self.status_codes_to_retry = status_codes_to_retry + self.number_of_tries = 0 + logger.info("Retry handler initialised with {}/{}/{}/{}/{}/".format( + self.backoff_type, + self.initial_delay, + self.maximum_number_of_retries, + self.maximum_duration_of_retry, + self.status_codes_to_retry + )) + + def should_retry(self, response): + logger.debug("Should retry?") + if response is None: + return True + if isinstance(response, requests.Response): + logger.info("is response") + status_code = str(response.status_code) + logger.info("status_code={}".format(status_code)) + if status_code in self.status_codes_to_retry: + logger.warning("HTTP error {}. Retrying.".format(status_code)) + self._compute_next_delay() + if self._is_next_delay_too_long(): + logger.info("_is_next_delay_too_long: should not.") + return False + if self._too_many_retries(): + logger.info("_too_many_retries: should not.") + return False + self._sleep() + return True + return False + + def recreate(self): + return RetryHandler( + backoff_type=self.backoff_type, + initial_delay=self.initial_delay, + maximum_number_of_retries=self.maximum_number_of_retries, + maximum_duration_of_retry=self.maximum_duration_of_retry, + status_codes_to_retry=self.status_codes_to_retry + ) + + def _compute_next_delay(self): + self.number_of_tries += 1 + if self.next_delay is None: + self.next_delay = self.initial_delay + return + if self.backoff_type=="linear": + # delay is same as last try + return + if self.backoff_type=="exponential": + self.next_delay = self.next_delay * 2 + + def _sleep(self): + if isinstance(self.next_delay, int): + logger.warning("Sleeping for {}s".format(self.next_delay)) + time.sleep(self.next_delay) + + def _too_many_retries(self): + if self.maximum_number_of_retries is None: + return False + if self.number_of_tries > self.maximum_number_of_retries: + logger.warning("Maximum number of retries reached. Not retrying.") + return True + return False + + def _is_next_delay_too_long(self): + if self.maximum_duration_of_retry is None: + return False + if self.next_delay >= self.maximum_duration_of_retry: + logger.warning("Sleep time before retry reached the max. Not retrying.") + return True + return False diff --git a/tests/python/integration/test_scenario.py b/tests/python/integration/test_scenario.py index 34d8bff..9603681 100644 --- a/tests/python/integration/test_scenario.py +++ b/tests/python/integration/test_scenario.py @@ -65,3 +65,10 @@ def test_run_api_connect_mtls(user_dss_clients): def test_run_api_connect_multipart_form_data(user_dss_clients): dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="MULTIPARTFORMDATA") + + +def test_run_api_connect_multipart_form_data_templating(user_dss_clients): + dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="MULTIPARTFORMDATATEMPLATING") + +def test_run_api_connect_retry_handler(user_dss_clients): + dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="RETRY") \ No newline at end of file diff --git a/tests/python/unit/test_common.py b/tests/python/unit/test_common.py index ddc6e6f..f9e95c5 100644 --- a/tests/python/unit/test_common.py +++ b/tests/python/unit/test_common.py @@ -1,5 +1,25 @@ from dku_utils import template_dict, join_url +from rest_api_client import RestAPIClient +from retry_handler import RetryHandler import pytest +import requests + + +class FakeSession: + def __init__(self, responses, request_times, clock): + self.responses = iter(responses) + self.request_times = request_times + self.clock = clock + + def request(self, *args, **kwargs): + self.request_times.append(self.clock[0]) + return next(self.responses) + + +def response_with_status(status_code): + response = requests.Response() + response.status_code = status_code + return response class TestCommonMethods: @@ -46,3 +66,72 @@ def test_join_url(self): assert "https://bla.com" == join_url("https://bla.com/", None) assert "https://bla.com" == join_url("https://bla.com", "") assert "https://bla.com" == join_url("https://bla.com", None) + + def run_retry_with_virtual_clock(self, monkeypatch, time_between_requests, backoff_type, initial_delay, + maximum_number_of_retries, response_codes): + clock = [0] + sleep_durations = [] + request_times = [] + + def sleep(duration): + sleep_durations.append(duration) + clock[0] += duration + + monkeypatch.setattr("rest_api_client.time.time", lambda: clock[0]) + monkeypatch.setattr("rest_api_client.time.sleep", sleep) + + client = object.__new__(RestAPIClient) + client.retry_handler = RetryHandler( + backoff_type=backoff_type, + initial_delay=initial_delay, + maximum_number_of_retries=maximum_number_of_retries, + status_codes_to_retry=["503"], + ) + client.time_between_requests = time_between_requests + client.time_last_request = None + client.session = FakeSession( + [response_with_status(status_code) for status_code in response_codes], + request_times, + clock, + ) + + client.request_with_errors_retry("GET", "https://example.test") + + return request_times, sleep_durations + + def test_exponential_retries_use_the_greater_of_backoff_and_throttle(self, monkeypatch): + request_times, _ = self.run_retry_with_virtual_clock( + monkeypatch, + time_between_requests=60, + backoff_type="exponential", + initial_delay=30, + maximum_number_of_retries=5, + response_codes=[503, 503, 503, 503, 503, 200], + ) + + assert request_times == [0, 60, 120, 240, 480, 960] + + def test_linear_retry_respects_throttling_rate_and_retry_budget(self, monkeypatch): + request_times, sleep_durations = self.run_retry_with_virtual_clock( + monkeypatch, + time_between_requests=30, + backoff_type="linear", + initial_delay=5, + maximum_number_of_retries=5, + response_codes=[503, 503, 503, 503, 503, 200], + ) + + assert request_times == [0, 30, 60, 90, 120, 150] + assert sleep_durations == [5, 25, 5, 25, 5, 25, 5, 25, 5, 25] + + def test_retry_does_not_make_another_request_after_exhausting_the_budget(self, monkeypatch): + request_times, _ = self.run_retry_with_virtual_clock( + monkeypatch, + time_between_requests=60, + backoff_type="linear", + initial_delay=1, + maximum_number_of_retries=5, + response_codes=[503, 503, 503, 503, 503, 503], + ) + + assert request_times == [0, 60, 120, 180, 240, 300]