Thanks to visit codestin.com
Credit goes to github.com

Skip to content

[ES-706907] Retry GetOperationStatus for http errors #145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 14, 2023
Merged

[ES-706907] Retry GetOperationStatus for http errors #145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0fbaa47
Retry GetOperationStatus for HTTPErrors too
Jun 13, 2023
0081610
Black the code
Jun 13, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions src/databricks/sql/thrift_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
import thrift.transport.TSocket
import thrift.transport.TTransport

import urllib3.exceptions

import databricks.sql.auth.thrift_http_client
from databricks.sql.auth.authenticators import AuthProvider
from databricks.sql.thrift_api.TCLIService import TCLIService, ttypes
Expand Down Expand Up @@ -324,6 +326,18 @@ def attempt_request(attempt):

logger.debug("Received response: {}".format(response))
return response

except urllib3.exceptions.HTTPError as err:
# retry on timeout. Happens a lot in Azure and it is safe as data has not been sent to server yet

gos_name = TCLIServiceClient.GetOperationStatus.__name__
if method.__name__ == gos_name:
retry_delay = bound_retry_delay(attempt, self._retry_delay_default)
logger.info(
f"GetOperationStatus failed with HTTP error and will be retried: {str(err)}"
)
else:
raise err
except OSError as err:
error = err
error_message = str(err)
Expand Down
48 changes: 47 additions & 1 deletion tests/unit/test_thrift_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from ssl import CERT_NONE, CERT_REQUIRED

import pyarrow
import urllib3

import databricks.sql
from databricks.sql.thrift_api.TCLIService import ttypes
Expand Down Expand Up @@ -1033,7 +1034,7 @@ 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.auth.thrift_http_client.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(
Expand Down Expand Up @@ -1089,6 +1090,51 @@ def test_make_request_will_retry_GetOperationStatus(
# The warnings should include this text
self.assertIn(f"{this_gos_name} failed with code {errno.EEXIST} and will attempt to retry", cm.output[0])

@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_for_http_error(
self, mock_retry_policy, mock_gos):

import urllib3.exceptions
mock_gos.side_effect = urllib3.exceptions.HTTPError("Read timed out")

import thrift, errno
from databricks.sql.thrift_api.TCLIService.TCLIService import Client
from databricks.sql.exc import RequestError
from databricks.sql.utils import NoRetryReason
from databricks.sql.auth.thrift_http_client import THttpClient

this_gos_name = "GetOperationStatus"
mock_gos.__name__ = this_gos_name

protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol(THttpClient)
client = Client(protocol)

req = ttypes.TGetOperationStatusReq(
operationHandle=self.operation_handle,
getProgressUpdate=False,
)

EXPECTED_RETRIES = 2

thrift_backend = ThriftBackend(
"foobar",
443,
"path", [],
auth_provider=AuthProvider(),
_retry_stop_after_attempts_count=EXPECTED_RETRIES,
_retry_delay_default=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):
Expand Down