|
6 | 6 | fromsslimportCERT_NONE,CERT_REQUIRED |
7 | 7 |
|
8 | 8 | importpyarrow |
| 9 | +importurllib3 |
9 | 10 |
|
10 | 11 | importdatabricks.sql |
11 | 12 | fromdatabricks.sql.thrift_api.TCLIServiceimportttypes |
@@ -1033,7 +1034,7 @@ def test_handle_execute_response_sets_active_op_handle(self): |
1033 | 1034 |
|
1034 | 1035 | self.assertEqual(mock_resp.operationHandle,mock_cursor.active_op_handle) |
1035 | 1036 |
|
1036 | | -@patch("thrift.transport.THttpClient.THttpClient") |
| 1037 | +@patch("databricks.sql.auth.thrift_http_client.THttpClient") |
1037 | 1038 | @patch("databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus") |
1038 | 1039 | @patch("databricks.sql.thrift_backend._retry_policy",new_callable=retry_policy_factory) |
1039 | 1040 | deftest_make_request_will_retry_GetOperationStatus( |
@@ -1089,6 +1090,51 @@ def test_make_request_will_retry_GetOperationStatus( |
1089 | 1090 | # The warnings should include this text |
1090 | 1091 | self.assertIn(f"{this_gos_name} failed with code{errno.EEXIST} and will attempt to retry",cm.output[0]) |
1091 | 1092 |
|
| 1093 | +@patch("databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus") |
| 1094 | +@patch("databricks.sql.thrift_backend._retry_policy",new_callable=retry_policy_factory) |
| 1095 | +deftest_make_request_will_retry_GetOperationStatus_for_http_error( |
| 1096 | +self,mock_retry_policy,mock_gos): |
| 1097 | + |
| 1098 | +importurllib3.exceptions |
| 1099 | +mock_gos.side_effect=urllib3.exceptions.HTTPError("Read timed out") |
| 1100 | + |
| 1101 | +importthrift,errno |
| 1102 | +fromdatabricks.sql.thrift_api.TCLIService.TCLIServiceimportClient |
| 1103 | +fromdatabricks.sql.excimportRequestError |
| 1104 | +fromdatabricks.sql.utilsimportNoRetryReason |
| 1105 | +fromdatabricks.sql.auth.thrift_http_clientimportTHttpClient |
| 1106 | + |
| 1107 | +this_gos_name="GetOperationStatus" |
| 1108 | +mock_gos.__name__=this_gos_name |
| 1109 | + |
| 1110 | +protocol=thrift.protocol.TBinaryProtocol.TBinaryProtocol(THttpClient) |
| 1111 | +client=Client(protocol) |
| 1112 | + |
| 1113 | +req=ttypes.TGetOperationStatusReq( |
| 1114 | +operationHandle=self.operation_handle, |
| 1115 | +getProgressUpdate=False, |
| 1116 | + ) |
| 1117 | + |
| 1118 | +EXPECTED_RETRIES=2 |
| 1119 | + |
| 1120 | +thrift_backend=ThriftBackend( |
| 1121 | +"foobar", |
| 1122 | +443, |
| 1123 | +"path", [], |
| 1124 | +auth_provider=AuthProvider(), |
| 1125 | +_retry_stop_after_attempts_count=EXPECTED_RETRIES, |
| 1126 | +_retry_delay_default=1) |
| 1127 | + |
| 1128 | + |
| 1129 | +withself.assertRaises(RequestError)ascm: |
| 1130 | +thrift_backend.make_request(client.GetOperationStatus,req) |
| 1131 | + |
| 1132 | + |
| 1133 | +self.assertEqual(NoRetryReason.OUT_OF_ATTEMPTS.value,cm.exception.context["no-retry-reason"]) |
| 1134 | +self.assertEqual(f'{EXPECTED_RETRIES}/{EXPECTED_RETRIES}',cm.exception.context["attempt"]) |
| 1135 | + |
| 1136 | + |
| 1137 | + |
1092 | 1138 |
|
1093 | 1139 | @patch("thrift.transport.THttpClient.THttpClient") |
1094 | 1140 | deftest_make_request_wont_retry_if_headers_not_present(self,t_transport_class): |
|
