Repository: apache/airflow

Issue:
Databricks hook: Retry also on HTTP Status 429 - rate limit exceeded
### Apache Airflow Provider(s)

databricks

### Versions of Apache Airflow Providers

2.2.0

### Apache Airflow version

2.2.3 (latest released)

### Operating System

Any

### Deployment

Other

### Deployment details

_No response_

### What happened

Operations aren't retried when Databricks API returns HTTP Status 429 - rate limit exceeded

### What you expected to happen

the operation should retry

### How to reproduce

this happens when you have multiple calls to API, especially when it happens outside of Airflow

### Anything else

_No response_

### Are you willing to submit PR?

- [X] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:

Resolving patch:
diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py
--- a/airflow/providers/databricks/hooks/databricks.py
+++ b/airflow/providers/databricks/hooks/databricks.py
@@ -100,6 +100,7 @@ class DatabricksHook(BaseDatabricksHook):
         service outages.
     :param retry_delay: The number of seconds to wait between retries (it
         might be a floating point number).
+    :param retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
     """
 
     hook_name = 'Databricks'
@@ -110,8 +111,9 @@ def __init__(
         timeout_seconds: int = 180,
         retry_limit: int = 3,
         retry_delay: float = 1.0,
+        retry_args: Optional[Dict[Any, Any]] = None,
     ) -> None:
-        super().__init__(databricks_conn_id, timeout_seconds, retry_limit, retry_delay)
+        super().__init__(databricks_conn_id, timeout_seconds, retry_limit, retry_delay, retry_args)
 
     def run_now(self, json: dict) -> int:
         """
diff --git a/airflow/providers/databricks/hooks/databricks_base.py b/airflow/providers/databricks/hooks/databricks_base.py
--- a/airflow/providers/databricks/hooks/databricks_base.py
+++ b/airflow/providers/databricks/hooks/databricks_base.py
@@ -22,15 +22,16 @@
 operators talk to the ``api/2.0/jobs/runs/submit``
 `endpoint <https://docs.databricks.com/api/latest/jobs.html#runs-submit>`_.
 """
+import copy
 import sys
 import time
-from time import sleep
 from typing import Any, Dict, Optional, Tuple
 from urllib.parse import urlparse
 
 import requests
 from requests import PreparedRequest, exceptions as requests_exceptions
 from requests.auth import AuthBase, HTTPBasicAuth
+from tenacity import RetryError, Retrying, retry_if_exception, stop_after_attempt, wait_exponential
 
 from airflow import __version__
 from airflow.exceptions import AirflowException
@@ -68,6 +69,7 @@ class BaseDatabricksHook(BaseHook):
         service outages.
     :param retry_delay: The number of seconds to wait between retries (it
         might be a floating point number).
+    :param retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
     """
 
     conn_name_attr = 'databricks_conn_id'
@@ -89,17 +91,33 @@ def __init__(
         timeout_seconds: int = 180,
         retry_limit: int = 3,
         retry_delay: float = 1.0,
+        retry_args: Optional[Dict[Any, Any]] = None,
     ) -> None:
         super().__init__()
         self.databricks_conn_id = databricks_conn_id
         self.timeout_seconds = timeout_seconds
         if retry_limit < 1:
-            raise ValueError('Retry limit must be greater than equal to 1')
+            raise ValueError('Retry limit must be greater than or equal to 1')
         self.retry_limit = retry_limit
         self.retry_delay = retry_delay
         self.aad_tokens: Dict[str, dict] = {}
         self.aad_timeout_seconds = 10
 
+        def my_after_func(retry_state):
+            self._log_request_error(retry_state.attempt_number, retry_state.outcome)
+
+        if retry_args:
+            self.retry_args = copy.copy(retry_args)
+            self.retry_args['retry'] = retry_if_exception(self._retryable_error)
+            self.retry_args['after'] = my_after_func
+        else:
+            self.retry_args = dict(
+                stop=stop_after_attempt(self.retry_limit),
+                wait=wait_exponential(min=self.retry_delay, max=(2 ** retry_limit)),
+                retry=retry_if_exception(self._retryable_error),
+                after=my_after_func,
+            )
+
     @cached_property
     def databricks_conn(self) -> Connection:
         return self.get_connection(self.databricks_conn_id)
@@ -143,6 +161,13 @@ def _parse_host(host: str) -> str:
             # In this case, host = xx.cloud.databricks.com
             return host
 
+    def _get_retry_object(self) -> Retrying:
+        """
+        Instantiates a retry object
+        :return: instance of Retrying class
+        """
+        return Retrying(**self.retry_args)
+
     def _get_aad_token(self, resource: str) -> str:
         """
         Function to get AAD token for given resource. Supports managed identity or service principal auth
@@ -154,60 +179,59 @@ def _get_aad_token(self, resource: str) -> str:
             return aad_token['token']
 
         self.log.info('Existing AAD token is expired, or going to expire soon. Refreshing...')
-        attempt_num = 1
-        while True:
-            try:
-                if self.databricks_conn.extra_dejson.get('use_azure_managed_identity', False):
-                    params = {
-                        "api-version": "2018-02-01",
-                        "resource": resource,
-                    }
-                    resp = requests.get(
-                        AZURE_METADATA_SERVICE_TOKEN_URL,
-                        params=params,
-                        headers={**USER_AGENT_HEADER, "Metadata": "true"},
-                        timeout=self.aad_timeout_seconds,
-                    )
-                else:
-                    tenant_id = self.databricks_conn.extra_dejson['azure_tenant_id']
-                    data = {
-                        "grant_type": "client_credentials",
-                        "client_id": self.databricks_conn.login,
-                        "resource": resource,
-                        "client_secret": self.databricks_conn.password,
-                    }
-                    azure_ad_endpoint = self.databricks_conn.extra_dejson.get(
-                        "azure_ad_endpoint", AZURE_DEFAULT_AD_ENDPOINT
-                    )
-                    resp = requests.post(
-                        AZURE_TOKEN_SERVICE_URL.format(azure_ad_endpoint, tenant_id),
-                        data=data,
-                        headers={**USER_AGENT_HEADER, 'Content-Type': 'application/x-www-form-urlencoded'},
-                        timeout=self.aad_timeout_seconds,
-                    )
-
-                resp.raise_for_status()
-                jsn = resp.json()
-                if 'access_token' not in jsn or jsn.get('token_type') != 'Bearer' or 'expires_on' not in jsn:
-                    raise AirflowException(f"Can't get necessary data from AAD token: {jsn}")
-
-                token = jsn['access_token']
-                self.aad_tokens[resource] = {'token': token, 'expires_on': int(jsn["expires_on"])}
-
-                return token
-            except requests_exceptions.RequestException as e:
-                if not self._retryable_error(e):
-                    raise AirflowException(
-                        f'Response: {e.response.content}, Status Code: {e.response.status_code}'
-                    )
-
-                self._log_request_error(attempt_num, e.strerror)
-
-            if attempt_num == self.retry_limit:
-                raise AirflowException(f'API requests to Azure failed {self.retry_limit} times. Giving up.')
-
-            attempt_num += 1
-            sleep(self.retry_delay)
+        try:
+            for attempt in self._get_retry_object():
+                with attempt:
+                    if self.databricks_conn.extra_dejson.get('use_azure_managed_identity', False):
+                        params = {
+                            "api-version": "2018-02-01",
+                            "resource": resource,
+                        }
+                        resp = requests.get(
+                            AZURE_METADATA_SERVICE_TOKEN_URL,
+                            params=params,
+                            headers={**USER_AGENT_HEADER, "Metadata": "true"},
+                            timeout=self.aad_timeout_seconds,
+                        )
+                    else:
+                        tenant_id = self.databricks_conn.extra_dejson['azure_tenant_id']
+                        data = {
+                            "grant_type": "client_credentials",
+                            "client_id": self.databricks_conn.login,
+                            "resource": resource,
+                            "client_secret": self.databricks_conn.password,
+                        }
+                        azure_ad_endpoint = self.databricks_conn.extra_dejson.get(
+                            "azure_ad_endpoint", AZURE_DEFAULT_AD_ENDPOINT
+                        )
+                        resp = requests.post(
+                            AZURE_TOKEN_SERVICE_URL.format(azure_ad_endpoint, tenant_id),
+                            data=data,
+                            headers={
+                                **USER_AGENT_HEADER,
+                                'Content-Type': 'application/x-www-form-urlencoded',
+                            },
+                            timeout=self.aad_timeout_seconds,
+                        )
+
+                    resp.raise_for_status()
+                    jsn = resp.json()
+                    if (
+                        'access_token' not in jsn
+                        or jsn.get('token_type') != 'Bearer'
+                        or 'expires_on' not in jsn
+                    ):
+                        raise AirflowException(f"Can't get necessary data from AAD token: {jsn}")
+
+                    token = jsn['access_token']
+                    self.aad_tokens[resource] = {'token': token, 'expires_on': int(jsn["expires_on"])}
+                    break
+        except RetryError:
+            raise AirflowException(f'API requests to Azure failed {self.retry_limit} times. Giving up.')
+        except requests_exceptions.HTTPError as e:
+            raise AirflowException(f'Response: {e.response.content}, Status Code: {e.response.status_code}')
+
+        return token
 
     def _get_aad_headers(self) -> dict:
         """
@@ -279,14 +303,6 @@ def _get_token(self, raise_error: bool = False) -> Optional[str]:
 
         return None
 
-    @staticmethod
-    def _retryable_error(exception) -> bool:
-        return (
-            isinstance(exception, (requests_exceptions.ConnectionError, requests_exceptions.Timeout))
-            or exception.response is not None
-            and exception.response.status_code >= 500
-        )
-
     def _log_request_error(self, attempt_num: int, error: str) -> None:
         self.log.error('Attempt %s API Request to Databricks failed with reason: %s', attempt_num, error)
 
@@ -327,36 +343,32 @@ def _do_api_call(self, endpoint_info: Tuple[str, str], json: Optional[Dict[str,
         else:
             raise AirflowException('Unexpected HTTP Method: ' + method)
 
-        attempt_num = 1
-        while True:
-            try:
-                response = request_func(
-                    url,
-                    json=json if method in ('POST', 'PATCH') else None,
-                    params=json if method == 'GET' else None,
-                    auth=auth,
-                    headers=headers,
-                    timeout=self.timeout_seconds,
-                )
-                response.raise_for_status()
-                return response.json()
-            except requests_exceptions.RequestException as e:
-                if not self._retryable_error(e):
-                    # In this case, the user probably made a mistake.
-                    # Don't retry.
-                    raise AirflowException(
-                        f'Response: {e.response.content}, Status Code: {e.response.status_code}'
+        try:
+            for attempt in self._get_retry_object():
+                with attempt:
+                    response = request_func(
+                        url,
+                        json=json if method in ('POST', 'PATCH') else None,
+                        params=json if method == 'GET' else None,
+                        auth=auth,
+                        headers=headers,
+                        timeout=self.timeout_seconds,
                     )
+                    response.raise_for_status()
+                    return response.json()
+        except RetryError:
+            raise AirflowException(f'API requests to Databricks failed {self.retry_limit} times. Giving up.')
+        except requests_exceptions.HTTPError as e:
+            raise AirflowException(f'Response: {e.response.content}, Status Code: {e.response.status_code}')
 
-                self._log_request_error(attempt_num, str(e))
-
-            if attempt_num == self.retry_limit:
-                raise AirflowException(
-                    f'API requests to Databricks failed {self.retry_limit} times. Giving up.'
-                )
-
-            attempt_num += 1
-            sleep(self.retry_delay)
+    @staticmethod
+    def _retryable_error(exception: BaseException) -> bool:
+        if not isinstance(exception, requests_exceptions.RequestException):
+            return False
+        return isinstance(exception, (requests_exceptions.ConnectionError, requests_exceptions.Timeout)) or (
+            exception.response is not None
+            and (exception.response.status_code >= 500 or exception.response.status_code == 429)
+        )
 
 
 class _TokenAuth(AuthBase):
diff --git a/airflow/providers/databricks/operators/databricks.py b/airflow/providers/databricks/operators/databricks.py
--- a/airflow/providers/databricks/operators/databricks.py
+++ b/airflow/providers/databricks/operators/databricks.py
@@ -245,6 +245,7 @@ class DatabricksSubmitRunOperator(BaseOperator):
         unreachable. Its value must be greater than or equal to 1.
     :param databricks_retry_delay: Number of seconds to wait between retries (it
             might be a floating point number).
+    :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
     :param do_xcom_push: Whether we should push run_id and run_page_url to xcom.
     """
 
@@ -274,6 +275,7 @@ def __init__(
         polling_period_seconds: int = 30,
         databricks_retry_limit: int = 3,
         databricks_retry_delay: int = 1,
+        databricks_retry_args: Optional[Dict[Any, Any]] = None,
         do_xcom_push: bool = False,
         idempotency_token: Optional[str] = None,
         access_control_list: Optional[List[Dict[str, str]]] = None,
@@ -287,6 +289,7 @@ def __init__(
         self.polling_period_seconds = polling_period_seconds
         self.databricks_retry_limit = databricks_retry_limit
         self.databricks_retry_delay = databricks_retry_delay
+        self.databricks_retry_args = databricks_retry_args
         self.wait_for_termination = wait_for_termination
         if tasks is not None:
             self.json['tasks'] = tasks
@@ -327,6 +330,7 @@ def _get_hook(self) -> DatabricksHook:
             self.databricks_conn_id,
             retry_limit=self.databricks_retry_limit,
             retry_delay=self.databricks_retry_delay,
+            retry_args=self.databricks_retry_args,
         )
 
     def execute(self, context: 'Context'):
@@ -484,6 +488,7 @@ class DatabricksRunNowOperator(BaseOperator):
         this run. By default the operator will poll every 30 seconds.
     :param databricks_retry_limit: Amount of times retry if the Databricks backend is
         unreachable. Its value must be greater than or equal to 1.
+    :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
     :param do_xcom_push: Whether we should push run_id and run_page_url to xcom.
     """
 
@@ -508,6 +513,7 @@ def __init__(
         polling_period_seconds: int = 30,
         databricks_retry_limit: int = 3,
         databricks_retry_delay: int = 1,
+        databricks_retry_args: Optional[Dict[Any, Any]] = None,
         do_xcom_push: bool = False,
         wait_for_termination: bool = True,
         **kwargs,
@@ -519,6 +525,7 @@ def __init__(
         self.polling_period_seconds = polling_period_seconds
         self.databricks_retry_limit = databricks_retry_limit
         self.databricks_retry_delay = databricks_retry_delay
+        self.databricks_retry_args = databricks_retry_args
         self.wait_for_termination = wait_for_termination
 
         if job_id is not None:
@@ -546,6 +553,7 @@ def _get_hook(self) -> DatabricksHook:
             self.databricks_conn_id,
             retry_limit=self.databricks_retry_limit,
             retry_delay=self.databricks_retry_delay,
+            retry_args=self.databricks_retry_args,
         )
 
     def execute(self, context: 'Context'):


Test patch:
diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py
--- a/tests/providers/databricks/hooks/test_databricks.py
+++ b/tests/providers/databricks/hooks/test_databricks.py
@@ -24,6 +24,7 @@
 from unittest import mock
 
 import pytest
+import tenacity
 from requests import exceptions as requests_exceptions
 from requests.auth import HTTPBasicAuth
 
@@ -50,6 +51,11 @@
 RUN_ID = 1
 JOB_ID = 42
 JOB_NAME = 'job-name'
+DEFAULT_RETRY_NUMBER = 3
+DEFAULT_RETRY_ARGS = dict(
+    wait=tenacity.wait_none(),
+    stop=tenacity.stop_after_attempt(DEFAULT_RETRY_NUMBER),
+)
 HOST = 'xx.cloud.databricks.com'
 HOST_WITH_SCHEME = 'https://xx.cloud.databricks.com'
 LOGIN = 'login'
@@ -228,6 +234,8 @@ def test_init_bad_retry_limit(self):
             DatabricksHook(retry_limit=0)
 
     def test_do_api_call_retries_with_retryable_error(self):
+        hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
+
         for exception in [
             requests_exceptions.ConnectionError,
             requests_exceptions.SSLError,
@@ -236,25 +244,41 @@ def test_do_api_call_retries_with_retryable_error(self):
             requests_exceptions.HTTPError,
         ]:
             with mock.patch('airflow.providers.databricks.hooks.databricks_base.requests') as mock_requests:
-                with mock.patch.object(self.hook.log, 'error') as mock_errors:
+                with mock.patch.object(hook.log, 'error') as mock_errors:
                     setup_mock_requests(mock_requests, exception)
 
                     with pytest.raises(AirflowException):
-                        self.hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
+                        hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
+
+                    assert mock_errors.call_count == DEFAULT_RETRY_NUMBER
+
+    def test_do_api_call_retries_with_too_many_requests(self):
+        hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
+
+        with mock.patch('airflow.providers.databricks.hooks.databricks_base.requests') as mock_requests:
+            with mock.patch.object(hook.log, 'error') as mock_errors:
+                setup_mock_requests(mock_requests, requests_exceptions.HTTPError, status_code=429)
+
+                with pytest.raises(AirflowException):
+                    hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
 
-                    assert mock_errors.call_count == self.hook.retry_limit
+                assert mock_errors.call_count == DEFAULT_RETRY_NUMBER
 
     @mock.patch('airflow.providers.databricks.hooks.databricks_base.requests')
     def test_do_api_call_does_not_retry_with_non_retryable_error(self, mock_requests):
+        hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
+
         setup_mock_requests(mock_requests, requests_exceptions.HTTPError, status_code=400)
 
-        with mock.patch.object(self.hook.log, 'error') as mock_errors:
+        with mock.patch.object(hook.log, 'error') as mock_errors:
             with pytest.raises(AirflowException):
-                self.hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
+                hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
 
             mock_errors.assert_not_called()
 
     def test_do_api_call_succeeds_after_retrying(self):
+        hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
+
         for exception in [
             requests_exceptions.ConnectionError,
             requests_exceptions.SSLError,
@@ -263,20 +287,18 @@ def test_do_api_call_succeeds_after_retrying(self):
             requests_exceptions.HTTPError,
         ]:
             with mock.patch('airflow.providers.databricks.hooks.databricks_base.requests') as mock_requests:
-                with mock.patch.object(self.hook.log, 'error') as mock_errors:
+                with mock.patch.object(hook.log, 'error') as mock_errors:
                     setup_mock_requests(
                         mock_requests, exception, error_count=2, response_content={'run_id': '1'}
                     )
 
-                    response = self.hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
+                    response = hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
 
                     assert mock_errors.call_count == 2
                     assert response == {'run_id': '1'}
 
-    @mock.patch('airflow.providers.databricks.hooks.databricks_base.sleep')
-    def test_do_api_call_waits_between_retries(self, mock_sleep):
-        retry_delay = 5
-        self.hook = DatabricksHook(retry_delay=retry_delay)
+    def test_do_api_call_custom_retry(self):
+        hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
 
         for exception in [
             requests_exceptions.ConnectionError,
@@ -286,16 +308,13 @@ def test_do_api_call_waits_between_retries(self, mock_sleep):
             requests_exceptions.HTTPError,
         ]:
             with mock.patch('airflow.providers.databricks.hooks.databricks_base.requests') as mock_requests:
-                with mock.patch.object(self.hook.log, 'error'):
-                    mock_sleep.reset_mock()
+                with mock.patch.object(hook.log, 'error') as mock_errors:
                     setup_mock_requests(mock_requests, exception)
 
                     with pytest.raises(AirflowException):
-                        self.hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
+                        hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
 
-                    assert len(mock_sleep.mock_calls) == self.hook.retry_limit - 1
-                    calls = [mock.call(retry_delay), mock.call(retry_delay)]
-                    mock_sleep.assert_has_calls(calls)
+                    assert mock_errors.call_count == DEFAULT_RETRY_NUMBER
 
     @mock.patch('airflow.providers.databricks.hooks.databricks_base.requests')
     def test_do_api_call_patch(self, mock_requests):
@@ -796,7 +815,7 @@ def setUp(self, session=None):
             }
         )
         session.commit()
-        self.hook = DatabricksHook()
+        self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
 
     @mock.patch('airflow.providers.databricks.hooks.databricks_base.requests')
     def test_submit_run(self, mock_requests):
@@ -838,7 +857,7 @@ def setUp(self, session=None):
             }
         )
         session.commit()
-        self.hook = DatabricksHook()
+        self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
 
     @mock.patch('airflow.providers.databricks.hooks.databricks_base.requests')
     def test_submit_run(self, mock_requests):
@@ -883,7 +902,7 @@ def setUp(self, session=None):
             }
         )
         session.commit()
-        self.hook = DatabricksHook()
+        self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
 
     @mock.patch('airflow.providers.databricks.hooks.databricks_base.requests')
     def test_submit_run(self, mock_requests):
@@ -931,7 +950,7 @@ def setUp(self, session=None):
             }
         )
         session.commit()
-        self.hook = DatabricksHook()
+        self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
 
     @mock.patch('airflow.providers.databricks.hooks.databricks_base.requests')
     def test_submit_run(self, mock_requests):
diff --git a/tests/providers/databricks/operators/test_databricks.py b/tests/providers/databricks/operators/test_databricks.py
--- a/tests/providers/databricks/operators/test_databricks.py
+++ b/tests/providers/databricks/operators/test_databricks.py
@@ -224,7 +224,10 @@ def test_exec_success(self, db_mock_class):
             {'new_cluster': NEW_CLUSTER, 'notebook_task': NOTEBOOK_TASK, 'run_name': TASK_ID}
         )
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
 
         db_mock.submit_run.assert_called_once_with(expected)
@@ -257,7 +260,10 @@ def test_exec_failure(self, db_mock_class):
             }
         )
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
         db_mock.submit_run.assert_called_once_with(expected)
         db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
@@ -297,7 +303,10 @@ def test_wait_for_termination(self, db_mock_class):
             {'new_cluster': NEW_CLUSTER, 'notebook_task': NOTEBOOK_TASK, 'run_name': TASK_ID}
         )
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
 
         db_mock.submit_run.assert_called_once_with(expected)
@@ -322,7 +331,10 @@ def test_no_wait_for_termination(self, db_mock_class):
             {'new_cluster': NEW_CLUSTER, 'notebook_task': NOTEBOOK_TASK, 'run_name': TASK_ID}
         )
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
 
         db_mock.submit_run.assert_called_once_with(expected)
@@ -445,7 +457,10 @@ def test_exec_success(self, db_mock_class):
         )
 
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
         db_mock.run_now.assert_called_once_with(expected)
         db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
@@ -475,7 +490,10 @@ def test_exec_failure(self, db_mock_class):
             }
         )
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
         db_mock.run_now.assert_called_once_with(expected)
         db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
@@ -513,7 +531,10 @@ def test_wait_for_termination(self, db_mock_class):
             }
         )
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
 
         db_mock.run_now.assert_called_once_with(expected)
@@ -540,7 +561,10 @@ def test_no_wait_for_termination(self, db_mock_class):
             }
         )
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
 
         db_mock.run_now.assert_called_once_with(expected)
@@ -582,7 +606,10 @@ def test_exec_with_job_name(self, db_mock_class):
         )
 
         db_mock_class.assert_called_once_with(
-            DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
         )
         db_mock.find_job_id_by_name.assert_called_once_with(JOB_NAME)
         db_mock.run_now.assert_called_once_with(expected)