Retry attempts that fail due to a connection timeout

src/databricks/sql/thrift_backend.py

@@ -5,6 +5,7 @@
 import threading
 from uuid import uuid4
 from ssl import CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED, create_default_context
+import socket
 
 import pyarrow
 import thrift.transport.THttpClient
@@ -39,6 +40,7 @@
     "_retry_delay_max": (float, 60, 5, 3600),
     "_retry_stop_after_attempts_count": (int, 30, 1, 60),
     "_retry_stop_after_attempts_duration": (float, 900, 1, 86400),
+    "_retry_delay_default": (float, 5, 1, 60),
 }
 
 
@@ -79,6 +81,8 @@ def __init__(
         #   next calculated pre-retry delay would go past
         #   _retry_stop_after_attempts_duration, stop now.)
         #
+        # _retry_delay_default                   (default: 5)
+        #   used when Retry-After is not specified by the server
         # _retry_stop_after_attempts_count
         #  The maximum number of times we should retry retryable requests (defaults to 24)
         # _socket_timeout
@@ -243,7 +247,7 @@ def _handle_request_error(self, error_info, attempt, elapsed):
     # FUTURE: Consider moving to https://github.com/litl/backoff or
     # https://github.com/jd/tenacity for retry logic.
     def make_request(self, method, request):
-        """Execute given request, attempting retries when receiving HTTP 429/503.
+        """Execute given request, attempting retries when TCP connection fils or when receiving HTTP 429/503.
 
         For delay between attempts, honor the given Retry-After header, but with bounds.
         Use lower bound of expontial-backoff based on _retry_delay_min,
@@ -261,8 +265,12 @@ def get_elapsed():
             return time.time() - t0
 
         def extract_retry_delay(attempt):
-            # encapsulate retry checks, returns None || delay-in-secs
-            # Retry IFF 429/503 code + Retry-After header set
+            """
+            Encapsulate retry checks based on HTTP headers. Returns None || delay-in-secs
+
+            Retry IFF 429/503 code + Retry-After header set
+            """
+
             http_code = getattr(self._transport, "code", None)
             retry_after = getattr(self._transport, "headers", {}).get("Retry-After")
             if http_code in [429, 503] and retry_after:
@@ -279,24 +287,41 @@ def attempt_request(attempt):
             # - non-None method_return -> success, return and be done
             # - non-None retry_delay -> sleep delay before retry
             # - error, error_message always set when available
+
+            error = None
+
+            # If retry_delay is None the request is treated as non-retryable
+            retry_delay = None
             try:
                 logger.debug("Sending request: {}".format(request))
                 response = method(request)
                 logger.debug("Received response: {}".format(response))
                 return response
-            except Exception as error:
+            except socket.timeout as err:
+                # socket.timeout means the connection was timed out by the socket package
+                retry_delay = self._retry_delay_default
+                error_message = str(err)
+                error = err
+            except OSError as err:
+                # OSError 110 means the connection was timed out by the operating system
+                if "Errno 110" in str(err):
+                    retry_delay = self._retry_delay_default
+                error_message = str(err)
+                error = err
+            except Exception as err:
                 retry_delay = extract_retry_delay(attempt)
                 error_message = ThriftBackend._extract_error_message_from_headers(
                     getattr(self._transport, "headers", {})
                 )
-                return RequestErrorInfo(
-                    error=error,
-                    error_message=error_message,
-                    retry_delay=retry_delay,
-                    http_code=getattr(self._transport, "code", None),
-                    method=method.__name__,
-                    request=request,
-                )
+                error = err
+            return RequestErrorInfo(
+                error=error,
+                error_message=error_message,
+                retry_delay=retry_delay,
+                http_code=getattr(self._transport, "code", None),
+                method=method.__name__,
+                request=request,
+            )
 
         # The real work:
         # - for each available attempt:

tests/e2e/driver_tests.py

@@ -7,7 +7,7 @@
 import sys
 import threading
 import time
-from unittest import loader, skipIf, skipUnless, TestCase
+from unittest import loader, skipIf, skipUnless, TestCase, mock
 from uuid import uuid4
 
 import numpy as np
@@ -360,6 +360,42 @@ def execute_really_long_query():
             cursor.execute("SELECT * FROM range(3)")
             self.assertEqual(len(cursor.fetchall()), 3)
 
+    def test_retry_after_connection_timeout(self):
+
+        import socket
+        from databricks.sql.thrift_api.TCLIService import ttypes
+
+        ATTEMPTS_TO_TRY = 2
+
+        with self.cursor({}) as cursor:
+            cursor.execute("SELECT id FROM RANGE(10)")
+            op_handle = cursor.active_op_handle
+            thrift_backend = cursor.thrift_backend
+
+
+        thrift_backend._retry_stop_after_attempts_count = ATTEMPTS_TO_TRY
+        req = ttypes.TGetOperationStatusReq(
+            operationHandle=op_handle,
+            getProgressUpdate=False,
+        )
+
+
+        with self.assertRaises(OperationalError) as cm:
+            with mock.patch("socket.create_connection") as mock_create_connection:
+                mock_create_connection.side_effect = OSError("[Errno 110]: Connection timed out")
+                thrift_backend.make_request(thrift_backend._client.GetOperationStatus, req)
+
+            self.assertIn("{0}/{0}".format(ATTEMPTS_TO_TRY), cm.exception.message_with_context())
+
+        with self.assertRaises(OperationalError) as cm:
+            with mock.patch("socket.create_connection") as mock_create_connection:
+                mock_create_connection.side_effect = socket.timeout
+                thrift_backend.make_request(thrift_backend._client.GetOperationStatus, req)
+
+            self.assertIn("{0}/{0}".format(ATTEMPTS_TO_TRY), cm.exception.message_with_context())
+
+
+
     @skipIf(pysql_has_version('<', '2'), 'requires pysql v2')
     def test_can_execute_command_after_failure(self):
         with self.cursor({}) as cursor:

tests/unit/test_thrift_backend.py

@@ -19,6 +19,7 @@ def retry_policy_factory():
         "_retry_delay_max":                     (float, 60, None, None),
         "_retry_stop_after_attempts_count":     (int, 30, None, None),
         "_retry_stop_after_attempts_duration":  (float, 900, None, None),
+        "_retry_delay_default":                 (float, 5, 0, 10),
     }
 
 
@@ -984,6 +985,63 @@ def test_make_request_wont_retry_if_headers_not_present(self, t_transport_class)
 
         self.assertIn("This method fails", str(cm.exception.message_with_context()))
 
+    @patch("thrift.transport.THttpClient.THttpClient")
+    def test_will_retry_on_connection_timeout(self, t_transport_class):
+
+        import socket
+
+        mock_method = Mock()
+        mock_method.__name__ = "method name"
+        mock_method.side_effect = socket.timeout
+
+        thrift_backend = ThriftBackend(
+            "foobar",
+            443,
+            "path",
+            [],
+            _retry_stop_after_attempts_count=2,
+            _retry_delay_default=0.25
+        )
+
+        with self.assertRaises(OperationalError) as cm:
+            thrift_backend.make_request(mock_method, Mock())
+
+            self.assertIn("2/2", cm.exception.message_with_context())
+
+        mock_method = Mock()
+        mock_method.__name__ = "method name"
+        mock_method.side_effect = OSError("[Errno 110] Connection timed out")
+
+        with self.assertRaises(OperationalError) as cm:
+            thrift_backend.make_request(mock_method, Mock())
+
+            self.assertIn("2/2", cm.exception.message_with_context())
+
+    @patch("thrift.transport.THttpClient.THttpClient")
+    def test_will_not_retry_on_non_timeout_oserror(self, t_transport_class):
+
+
+
+        mock_method = Mock()
+        mock_method.__name__ = "method name"
+        mock_method.side_effect = OSError("I am not a timeout error")
+
+        thrift_backend = ThriftBackend(
+            "foobar",
+            443,
+            "path",
+            [],
+            _retry_stop_after_attempts_count=2,
+            _retry_delay_default=0.25
+        )
+
+        with self.assertRaises(OperationalError) as cm:
+            thrift_backend.make_request(mock_method, Mock())
+
+        self.assertIn("I am not a timeout error", str(cm.exception.message_with_context()))
+        self.assertIn("1/2", cm.exception.message_with_context())
+
+
     @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