Scope empty response retries to statement requests

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
 ])

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,