429, 503 Retries in Thrift Client

This PR introduces changes to retry if there is a `Retry-After` header and the response is 429 and 503

* New unit tests

Author: NiallEgan

cmdexec/clients/python/src/databricks/sql/client.py

@@ -86,6 +86,8 @@ def __init__(self,
         #   Set client SSL certificate.
         # _session_id
         #   Specify the session id of the connection. For Redash use only.
+        # _max_number_of_retries
+        #  The maximum number of times we should retry retriable requests (defaults to 25)
 
         self.host = server_hostname
         self.port = kwargs.get("_port", 443)
@@ -117,7 +119,8 @@ def __init__(self,
             _tls_client_cert_file=kwargs.get("_tls_client_cert_file"),
             _tls_client_cert_key_file=kwargs.get("_tls_client_cert_key_file"),
             _tls_client_cert_key_password=kwargs.get("_tls_client_cert_key_password"),
-            _connection_uri=kwargs.get("_connection_uri"))
+            _connection_uri=kwargs.get("_connection_uri"),
+            _max_number_of_retries=kwargs.get("_max_number_of_retries", 25))
 
         self._session_handle = self.thrift_backend.open_session(
             session_id=kwargs.get("_session_id"))

cmdexec/clients/python/src/databricks/sql/thrift_backend.py

@@ -1,4 +1,6 @@
 import logging
+import time
+import threading
 from uuid import uuid4
 from ssl import CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED, create_default_context
 
@@ -20,6 +22,7 @@ class ThriftBackend:
     CLOSED_OP_STATE = ttypes.TOperationState.CLOSED_STATE
     ERROR_OP_STATE = ttypes.TOperationState.ERROR_STATE
     BIT_MASKS = [1, 2, 4, 8, 16, 32, 64, 128]
+    ERROR_MSG_HEADER = "X-Thriftserver-Error-Message"
 
     def __init__(self, server_hostname: str, port, http_path: str, http_headers, **kwargs):
         # Internal arguments in **kwargs:
@@ -39,6 +42,8 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
         #   See https://docs.python.org/3/library/ssl.html#ssl.SSLContext.load_cert_chain
         # _connection_uri
         #   Overrides server_hostname and http_path.
+        # _max_number_of_retries
+        #  The maximum number of times we should retry retryable requests (defaults to 25)
 
         port = port or 443
         if kwargs.get("_connection_uri"):
@@ -49,6 +54,8 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
         else:
             raise ValueError("No valid connection settings.")
 
+        self._max_number_of_retries = kwargs.get("_max_number_of_retries", 25)
+
         # Configure tls context
         ssl_context = create_default_context(cafile=kwargs.get("_tls_trusted_ca_file"))
         if kwargs.get("_tls_no_verify") is True:
@@ -85,21 +92,53 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
             self._transport.close()
             raise
 
+        self._request_lock = threading.RLock()
+
     @staticmethod
     def _check_response_for_error(response):
         if response.status and response.status.statusCode in \
                 [ttypes.TStatusCode.ERROR_STATUS, ttypes.TStatusCode.INVALID_HANDLE_STATUS]:
             raise DatabaseError(response.status.errorMessage)
 
-    @staticmethod
-    def make_request(method, request):
+    def make_request(self, method, request, attempt_number=1):
         try:
+            # We have a lock here because .cancel can be called from a separate thread.
+            # We do not want threads to be simultaneously sharing the Thrift Transport
+            # because we use its state to determine retries
+            self._request_lock.acquire()
             response = method(request)
             logger.debug("Received response: {}".format(response))
             ThriftBackend._check_response_for_error(response)
             return response
-        except TException as error:
-            raise OperationalError("Error during Thrift request", error)
+        except Exception as error:
+            # _transport.code isn't necessarily set :(
+            code_and_headers_is_set = hasattr(self._transport, 'code') \
+                                      and hasattr(self._transport, 'headers')
+            # We only retry if a Retry-After header is set
+            if code_and_headers_is_set and self._transport.code in [503, 429] and \
+                    "Retry-After" in self._transport.headers and \
+                    attempt_number <= self._max_number_of_retries:
+                retry_time_seconds = int(self._transport.headers["Retry-After"])
+                if self.ERROR_MSG_HEADER in self._transport.headers:
+                    error_message = self._transport.headers[self.ERROR_MSG_HEADER]
+                else:
+                    error_message = str(error)
+                logger.warning("Received retryable error during {}. Request: {} Error: {}".format(
+                    method, request, error_message))
+                logger.warning("Retrying in {} seconds. This is attempt number {}".format(
+                    retry_time_seconds, attempt_number))
+                time.sleep(retry_time_seconds)
+                return self.make_request(method, request, attempt_number + 1)
+            else:
+                logger.error("Received error when issuing: {}".format(request))
+                if hasattr(self._transport, "headers") and \
+                   self.ERROR_MSG_HEADER in self._transport.headers:
+                    error_message = self._transport.headers[self.ERROR_MSG_HEADER]
+                    raise OperationalError("Error during Thrift request: {}".format(error_message))
+                else:
+                    raise OperationalError("Error during Thrift request", error)
+        finally:
+            self._request_lock.release()
 
     def _check_protocol_version(self, t_open_session_resp):
         protocol_version = t_open_session_resp.serverProtocolVersion

cmdexec/clients/python/tests/test_thrift_backend.py

@@ -43,8 +43,9 @@ def test_make_request_checks_thrift_status_code(self):
         mock_response = Mock()
         mock_response.status.statusCode = ttypes.TStatusCode.ERROR_STATUS
         mock_method = lambda _: mock_response
+        thrift_backend = ThriftBackend("foo", 123, "bar", [])
         with self.assertRaises(DatabaseError):
-            ThriftBackend.make_request(mock_method, Mock())
+            thrift_backend.make_request(mock_method, Mock())
 
     def _make_type_desc(self, type):
         return ttypes.TTypeDesc(types=[ttypes.TTypeEntry(ttypes.TPrimitiveTypeEntry(type=type))])
@@ -174,22 +175,25 @@ def test_hive_schema_to_description_preserves_column_names_and_types(self):
 
     def test_make_request_checks_status_code(self):
         error_codes = [ttypes.TStatusCode.ERROR_STATUS, ttypes.TStatusCode.INVALID_HANDLE_STATUS]
+        thrift_backend = ThriftBackend("foo", 123, "bar", [])
+
         for code in error_codes:
             mock_error_response = Mock()
             mock_error_response.status.statusCode = code
             mock_error_response.status.errorMessage = "a detailed error message"
             with self.assertRaises(DatabaseError) as cm:
-                ThriftBackend.make_request(lambda _: mock_error_response, Mock())
+                thrift_backend.make_request(lambda _: mock_error_response, Mock())
             self.assertIn("a detailed error message", str(cm.exception))
 
         success_codes = [
             ttypes.TStatusCode.SUCCESS_STATUS, ttypes.TStatusCode.SUCCESS_WITH_INFO_STATUS,
             ttypes.TStatusCode.STILL_EXECUTING_STATUS
         ]
+
         for code in success_codes:
             mock_response = Mock()
             mock_response.status.statusCode = code
-            ThriftBackend.make_request(lambda _: mock_response, Mock())
+            thrift_backend.make_request(lambda _: mock_response, Mock())
 
     def test_handle_execute_response_checks_operation_state_in_direct_results(self):
         for resp_type in self.execute_response_types:
@@ -746,6 +750,74 @@ def test_handle_execute_response_sets_active_op_handle(self):
 
         self.assertEqual(mock_resp.operationHandle, mock_cursor.active_op_handle)
 
+    @patch("thrift.transport.THttpClient.THttpClient")
+    def test_make_request_wont_retry_if_headers_not_present(self, t_transport_class):
+        t_transport_instance = t_transport_class.return_value
+        t_transport_instance.code = 429
+        t_transport_instance.headers = {"foo": "bar"}
+        mock_method = Mock()
+        mock_method.side_effect = Exception("This method fails")
+
+        thrift_backend = ThriftBackend("foobar", 443, "path", [])
+
+        with self.assertRaises(OperationalError) as cm:
+            thrift_backend.make_request(mock_method, Mock())
+
+        self.assertIn("This method fails", str(cm.exception))
+
+    @patch("thrift.transport.THttpClient.THttpClient")
+    def test_make_request_wont_retry_if_error_code_not_429_or_503(self, t_transport_class):
+        t_transport_instance = t_transport_class.return_value
+        t_transport_instance.code = 430
+        t_transport_instance.headers = {"Retry-After": "1"}
+        mock_method = Mock()
+        mock_method.side_effect = Exception("This method fails")
+
+        thrift_backend = ThriftBackend("foobar", 443, "path", [])
+
+        with self.assertRaises(OperationalError) as cm:
+            thrift_backend.make_request(mock_method, Mock())
+
+        self.assertIn("This method fails", str(cm.exception))
+
+    @patch("thrift.transport.THttpClient.THttpClient")
+    def test_make_request_will_retry_max_number_of_retries_times_if_retryable(
+            self, t_transport_class):
+        t_transport_instance = t_transport_class.return_value
+        t_transport_instance.code = 429
+        t_transport_instance.headers = {"Retry-After": "0"}
+        mock_method = Mock()
+        mock_method.side_effect = Exception("This method fails")
+
+        thrift_backend = ThriftBackend("foobar", 443, "path", [], _max_number_of_retries=13)
+
+        with self.assertRaises(OperationalError) as cm:
+            thrift_backend.make_request(mock_method, Mock())
+
+        self.assertIn("This method fails", str(cm.exception))
+
+        self.assertEqual(mock_method.call_count, 13 + 1)
+
+    @patch("thrift.transport.THttpClient.THttpClient")
+    def test_make_request_will_read_X_Thriftserver_Error_Message_if_set(self, t_transport_class):
+        t_transport_instance = t_transport_class.return_value
+        t_transport_instance.code = 429
+        t_transport_instance.headers = {
+            "Retry-After": "0",
+            "X-Thriftserver-Error-Message": "message2"
+        }
+        mock_method = Mock()
+        mock_method.side_effect = Exception("This method fails")
+
+        thrift_backend = ThriftBackend("foobar", 443, "path", [], _max_number_of_retries=13)
+
+        with self.assertRaises(OperationalError) as cm:
+            thrift_backend.make_request(mock_method, Mock())
+
+        self.assertIn("message2", str(cm.exception))
+
+        self.assertEqual(mock_method.call_count, 13 + 1)
+
 
 if __name__ == '__main__':
     unittest.main()

cmdexec/clients/python/tests/tests.py

@@ -325,6 +325,12 @@ def test_cancel_command_will_issue_warning_for_cancel_with_no_executing_command(
         self.assertTrue(logger_instance.warning.called)
         self.assertFalse(mock_thrift_backend.cancel_command.called)
 
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_max_number_of_retries_passthrough(self, mock_client_class):
+        databricks.sql.connect(_max_number_of_retries=53, **self.DUMMY_CONNECTION_ARGS)
+
+        self.assertEqual(mock_client_class.call_args[1]["_max_number_of_retries"], 53)
+
 
 class ResultSetTests(unittest.TestCase):
     def test_parse_type_converts_decimal(self):