diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index c5d33dbd..0980cda5 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1142,6 +1142,27 @@ def test_empty_200_response_retry(monkeypatch): assert get_retry.retry_count == attempts +@pytest.mark.parametrize("method_name", ["_get", "_post", "_delete", "_head"]) +def test_empty_200_transport_response_does_not_retry(method_name, monkeypatch): + http_resp = TrinoRequest.http.Response() + http_resp.status_code = 200 + http_resp._content = b"" + + retry = RetryRecorder(result=http_resp) + session_method = method_name.removeprefix("_") + monkeypatch.setattr(TrinoRequest.http.Session, session_method, retry) + + req = TrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession(user="test"), + max_attempts=3, + ) + + getattr(req, method_name)("URL") + assert retry.retry_count == 1 + + @pytest.mark.parametrize("status_code", [ 501 ]) diff --git a/trino/client.py b/trino/client.py index 9d4956bf..f8820534 100644 --- a/trino/client.py +++ b/trino/client.py @@ -638,11 +638,23 @@ def max_attempts(self, value: int) -> None: if value == 1: # No retry self._get = self._http_session.get self._post = self._http_session.post + self._get_statement = self._http_session.get + self._post_statement = self._http_session.post self._delete = self._http_session.delete self._head = self._http_session.head return with_retry = _retry_with( + self._handle_retry, + handled_exceptions=self._exceptions, + conditions=( + # need retry when there is no exception but the status code is 429, 502, 503, or 504 + lambda response: getattr(response, "status_code", None) + in (429, 502, 503, 504), + ), + max_attempts=self._max_attempts, + ) + with_statement_retry = _retry_with( self._handle_retry, handled_exceptions=self._exceptions, conditions=( @@ -651,12 +663,14 @@ def max_attempts(self, value: int) -> None: in (429, 502, 503, 504), # need retry when the server returns 200 with an empty body (transient under load) lambda response: getattr(response, "status_code", None) == 200 - and not getattr(response, "text", "").strip(), + and not getattr(response, "content", b""), ), max_attempts=self._max_attempts, ) self._get = with_retry(self._http_session.get) self._post = with_retry(self._http_session.post) + self._get_statement = with_statement_retry(self._http_session.get) + self._post_statement = with_statement_retry(self._http_session.post) self._delete = with_retry(self._http_session.delete) self._head = with_retry(self._http_session.head) @@ -686,7 +700,7 @@ def post(self, sql: str, additional_http_headers: Optional[Dict[str, Any]] = Non # explicitly to match the Trino JDBC client. Users may still override it. http_headers.setdefault(constants.HEADER_CONTENT_TYPE, constants.CONTENT_TYPE_TEXT_UTF8) - http_response = self._post( + http_response = self._post_statement( self.statement_url, data=data, headers=http_headers, @@ -696,7 +710,7 @@ def post(self, sql: str, additional_http_headers: Optional[Dict[str, Any]] = Non return http_response def get(self, url: str) -> Response: - return self._get( + return self._get_statement( url, headers=self.http_headers, timeout=self._request_timeout,