Skip to content

Commit

Permalink
Updated the retryable condition
Browse files Browse the repository at this point in the history
  • Loading branch information
jprakash-db committed Dec 20, 2024
1 parent fba9759 commit aefd44b
Show file tree
Hide file tree
Showing 4 changed files with 1 addition and 28 deletions.
11 changes: 0 additions & 11 deletions src/databricks/sql/auth/retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"
Expand Down
9 changes: 0 additions & 9 deletions src/databricks/sql/auth/thrift_http_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
)
3 changes: 1 addition & 2 deletions src/databricks/sql/thrift_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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):
Expand Down
6 changes: 0 additions & 6 deletions tests/unit/test_retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

0 comments on commit aefd44b

Please sign in to comment.