Retry GetOperationStatus if an OSError was raised during execution

Add retry_delay_default to use in this case.

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: susodapop

src/databricks/sql/thrift_backend.py

@@ -15,6 +15,9 @@
 
 from databricks.sql.thrift_api.TCLIService import TCLIService, ttypes
 from databricks.sql import *
+from databricks.sql.thrift_api.TCLIService.TCLIService import (
+    Client as TCLIServiceClient,
+)
 from databricks.sql.utils import (
     ArrowQueue,
     ExecuteResponse,
@@ -39,6 +42,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)
 }
 
 
@@ -71,6 +75,8 @@ def __init__(
         # _retry_delay_min                      (default: 1)
         # _retry_delay_max                      (default: 60)
         #   {min,max} pre-retry delay bounds
+        # _retry_delay_default                   (default: 5)
+        #   Only used when GetOperationStatus fails due to a TCP/OS Error. 
         # _retry_stop_after_attempts_count      (default: 30)
         #   total max attempts during retry sequence
         # _retry_stop_after_attempts_duration   (default: 900)
@@ -291,6 +297,13 @@ def attempt_request(attempt):
                 response = method(request)
                 logger.debug("Received response: {}".format(response))
                 return response
+            except OSError as err:
+                error = err
+                error_message = str(err)
+
+                gos_name = TCLIServiceClient.GetOperationStatus.__name__
+                if method.__name__ == gos_name:
+                    retry_delay = bound_retry_delay(attempt, self._retry_delay_default)
             except Exception as err:
                 error = err
                 retry_delay = extract_retry_delay(attempt)

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, 1, 60)
     }
 
 
@@ -968,6 +969,44 @@ 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")
+    @patch("databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus")
+    @patch("databricks.sql.thrift_backend._retry_policy", new_callable=retry_policy_factory)
+    def test_make_request_will_retry_GetOperationStatus(
+            self, mock_retry_policy, mock_GetOperationStatus, t_transport_class):
+       
+        import thrift
+        from databricks.sql.thrift_api.TCLIService.TCLIService import Client
+        from databricks.sql.exc import RequestError
+        from databricks.sql.utils import NoRetryReason
+
+        mock_GetOperationStatus.__name__ = "GetOperationStatus"
+        mock_GetOperationStatus.side_effect = TimeoutError(110)
+
+        protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol(t_transport_class)
+        client = Client(protocol)
+
+        req = ttypes.TGetOperationStatusReq(
+            operationHandle=self.operation_handle,
+            getProgressUpdate=False,
+        )
+        
+        EXPECTED_RETRIES = 2
+
+        thrift_backend = ThriftBackend(
+            "foobar",
+            443,
+            "path", [],
+            _retry_stop_after_attempts_count=EXPECTED_RETRIES,
+            _retry_delay_default=0.1)
+
+        with self.assertRaises(RequestError) as cm:
+            thrift_backend.make_request(client.GetOperationStatus, req)
+
+        self.assertEqual(NoRetryReason.OUT_OF_ATTEMPTS.value, cm.exception.context["no-retry-reason"])
+        self.assertEqual(f'{EXPECTED_RETRIES}/{EXPECTED_RETRIES}', cm.exception.context["attempt"])    
+    
+    
     @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