Updated the retryable condition

Author: jprakash-db

src/databricks/sql/auth/retry.py

@@ -241,14 +241,6 @@ def command_type(self) -> Optional[CommandType]:
     def command_type(self, value: CommandType) -> None:
         self._command_type = value
 
-    @property
-    def is_retryable(self) -> bool:
-        return self._is_retryable
-
-    @is_retryable.setter
-    def is_retryable(self, value: bool) -> None:
-        self._is_retryable = value
-
     @property
     def delay_default(self) -> float:
         """Time in seconds the connector will wait between requests polling a GetOperationStatus Request
@@ -369,9 +361,6 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if status_code == 501:
             raise NonRecoverableNetworkError("Received code 501 from server.")
 
-        if self.is_retryable == False:
-            return False, "Request is not retryable"
-
         # Request failed and this method is not retryable. We only retry POST requests.
         if not self._is_method_retryable(method):
             return False, "Only POST requests are retried"

src/databricks/sql/auth/thrift_http_client.py

@@ -216,12 +216,3 @@ def set_retry_command_type(self, value: CommandType):
             logger.warning(
                 "DatabricksRetryPolicy is currently bypassed. The CommandType cannot be set."
             )
-
-    def set_is_retryable(self, retryable: bool):
-        """Pass the provided retryable flag to the retry policy"""
-        if isinstance(self.retry_policy, DatabricksRetryPolicy):
-            self.retry_policy.is_retryable = retryable
-        else:
-            logger.warning(
-                "DatabricksRetryPolicy is currently bypassed. The is_retryable flag cannot be set."
-            )

src/databricks/sql/thrift_backend.py

@@ -376,7 +376,6 @@ def attempt_request(attempt):
                 if self.enable_v3_retries:
                     this_command_type = CommandType.get(this_method_name)
                     self._transport.set_retry_command_type(this_command_type)
-                    self._transport.set_is_retryable(retryable)
                     self._transport.startRetryTimer()
 
                 response = method(request)
@@ -461,7 +460,7 @@ def attempt_request(attempt):
         #       return on success
         #       if available: bounded delay and retry
         #       if not: raise error
-        max_attempts = self._retry_stop_after_attempts_count
+        max_attempts = self._retry_stop_after_attempts_count if retryable else 1
 
         # use index-1 counting for logging/human consistency
         for attempt in range(1, max_attempts + 1):

tests/unit/test_retry.py

@@ -84,9 +84,3 @@ def test_sleep__retry_after_present(self, t_mock, retry_policy, error_history):
         retry_policy.history = [error_history, error_history, error_history]
         retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "3"}))
         t_mock.assert_called_with(3)
-
-    def test_not_retryable__fetch_results_orientation_fetch_next(self, retry_policy):
-        HTTP_STATUS_CODES = [200, 429, 503, 504]
-        retry_policy.is_retryable = False
-        for status_code in HTTP_STATUS_CODES:
-            assert not retry_policy.is_retry("METHOD_NAME", status_code=status_code)