[GitHub] andrewmchen commented on a change in pull request #3570: [AIRFLOW-2709] Improve error handling in Databricks hook

[GitBox <gi...@apache.org>]

andrewmchen commented on a change in pull request #3570: [AIRFLOW-2709] Improve error handling in Databricks hook
URL: https://github.com/apache/incubator-airflow/pull/3570#discussion_r211753269
 
 

 ##########
 File path: airflow/contrib/hooks/databricks_hook.py
 ##########
 @@ -117,29 +123,51 @@ def _do_api_call(self, endpoint_info, json):
         else:
             raise AirflowException('Unexpected HTTP Method: ' + method)
 
-        for attempt_num in range(1, self.retry_limit + 1):
+        attempt_num = 1
+        while True:
             try:
                 response = request_func(
                     url,
                     json=json,
                     auth=auth,
                     headers=USER_AGENT_HEADER,
                     timeout=self.timeout_seconds)
-                if response.status_code == requests.codes.ok:
-                    return response.json()
-                else:
+                response.raise_for_status()
+                return response.json()
+            except (requests_exceptions.ConnectionError,
+                    requests_exceptions.Timeout) as e:
+                self._log_request_error(attempt_num, e)
+            except requests_exceptions.HTTPError as e:
+                response = e.response
+                if not self._retriable_error(response):
                     # In this case, the user probably made a mistake.
                     # Don't retry.
                     raise AirflowException('Response: {0}, Status Code: {1}'.format(
                         response.content, response.status_code))
-            except (requests_exceptions.ConnectionError,
-                    requests_exceptions.Timeout) as e:
-                self.log.error(
-                    'Attempt %s API Request to Databricks failed with reason: %s',
-                    attempt_num, e
-                )
-        raise AirflowException(('API requests to Databricks failed {} times. ' +
-                               'Giving up.').format(self.retry_limit))
+
+                self._log_request_error(attempt_num, e)
+
+            if attempt_num == self.retry_limit:
+                raise AirflowException(('API requests to Databricks failed {} times. ' +
+                                        'Giving up.').format(self.retry_limit))
+
+            attempt_num += 1
+            sleep(self.retry_delay)
+
+    def _log_request_error(self, attempt_num, error):
+        self.log.error(
+            'Attempt %s API Request to Databricks failed with reason: %s',
+            attempt_num, error
+        )
+
+    @staticmethod
+    def _retriable_error(response):
+        try:
+            error_code = response.json().get('error_code')
+            return error_code == 'TEMPORARILY_UNAVAILABLE'
 
 Review comment:
   Thanks, this is definitely a known issue and can happen sometimes when the gateway proxy is not behaving correctly... That being said, this fix is definitely necessary and appreciated!

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services