Retries for 429 and 503 HTTP errors

Add retry support to clients during HTTP 429/503 scenarios, with sane defaults.

Author: juliuszsompolski

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

@@ -57,8 +57,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)
+        # _retry_stop_after_attempts_count
+        #  The maximum number of attempts during a request retry sequence (defaults to 24)
 
         self.host = server_hostname
         self.port = kwargs.get("_port", 443)
@@ -91,7 +91,7 @@ def __init__(self,
             _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"),
-            _max_number_of_retries=kwargs.get("_max_number_of_retries", 25))
+            _retry_stop_after_attempts_count=kwargs.get("_retry_stop_after_attempts_count", 24))
 
         self._session_handle = self.thrift_backend.open_session(
             session_id=kwargs.get("_session_id"))

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

@@ -43,8 +43,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)
+        # _retry_stop_after_attempts_count
+        #  The maximum number of times we should retry retryable requests (defaults to 24)
 
         port = port or 443
         if kwargs.get("_connection_uri"):
@@ -55,7 +55,7 @@ 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)
+        self._retry_stop_after_attempts_count = kwargs.get("_retry_stop_after_attempts_count", 24)
 
         # Configure tls context
         ssl_context = create_default_context(cafile=kwargs.get("_tls_trusted_ca_file"))
@@ -101,6 +101,9 @@ def _check_response_for_error(response):
                 [ttypes.TStatusCode.ERROR_STATUS, ttypes.TStatusCode.INVALID_HANDLE_STATUS]:
             raise DatabaseError(response.status.errorMessage)
 
+    # FUTURE: Consider moving to https://github.com/litl/backoff or
+    # https://github.com/jd/tenacity for retry logic. Otherwise, copy from
+    # v1 client.
     def make_request(self, method, request, attempt_number=1):
         try:
             # We have a lock here because .cancel can be called from a separate thread.
@@ -118,7 +121,7 @@ def make_request(self, method, request, attempt_number=1):
             # 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:
+                    attempt_number <= self._retry_stop_after_attempts_count - 1:
                 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]

cmdexec/clients/python/tests/test_thrift_backend.py

@@ -808,22 +808,23 @@ def test_make_request_wont_retry_if_error_code_not_429_or_503(self, t_transport_
         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(
+    def test_make_request_will_retry_stop_after_attempts_count_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)
+        thrift_backend = ThriftBackend(
+            "foobar", 443, "path", [], _retry_stop_after_attempts_count=14)
 
         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)
+        self.assertEqual(mock_method.call_count, 14)
 
     @patch("thrift.transport.THttpClient.THttpClient")
     def test_make_request_will_read_X_Thriftserver_Error_Message_if_set(self, t_transport_class):
@@ -836,14 +837,15 @@ def test_make_request_will_read_X_Thriftserver_Error_Message_if_set(self, t_tran
         mock_method = Mock()
         mock_method.side_effect = Exception("This method fails")
 
-        thrift_backend = ThriftBackend("foobar", 443, "path", [], _max_number_of_retries=13)
+        thrift_backend = ThriftBackend(
+            "foobar", 443, "path", [], _retry_stop_after_attempts_count=14)
 
         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)
+        self.assertEqual(mock_method.call_count, 14)
 
     @staticmethod
     def make_table_and_desc(height, n_decimal_cols, width, precision, scale, int_constant,

cmdexec/clients/python/tests/tests.py

@@ -327,9 +327,9 @@ def test_cancel_command_will_issue_warning_for_cancel_with_no_executing_command(
 
     @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)
+        databricks.sql.connect(_retry_stop_after_attempts_count=54, **self.DUMMY_CONNECTION_ARGS)
 
-        self.assertEqual(mock_client_class.call_args[1]["_max_number_of_retries"], 53)
+        self.assertEqual(mock_client_class.call_args[1]["_retry_stop_after_attempts_count"], 54)
 
 
 if __name__ == '__main__':