Skip to content

Commit 2caf38d

Browse files
more defensive parsing, allow more method types in urllib3
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent d67eb7b commit 2caf38d

File tree

2 files changed

+54
-26
lines changed

2 files changed

+54
-26
lines changed

src/databricks/sql/auth/retry.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -130,8 +130,8 @@ def __init__(
130130
allowed_methods=["POST"],
131131
status_forcelist=[429, 503, *self.force_dangerous_codes],
132132
)
133-
134-
urllib3_kwargs.update(**_urllib_kwargs_we_care_about)
133+
_urllib_kwargs_we_care_about.update(**urllib3_kwargs)
134+
urllib3_kwargs = _urllib_kwargs_we_care_about
135135

136136
super().__init__(
137137
**urllib3_kwargs,

src/databricks/sql/backend/sea/utils/http_client.py

Lines changed: 52 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,12 @@
1212
from databricks.sql.auth.authenticators import AuthProvider
1313
from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy
1414
from databricks.sql.types import SSLOptions
15-
from databricks.sql.exc import RequestError, MaxRetryDurationError
15+
from databricks.sql.exc import (
16+
RequestError,
17+
MaxRetryDurationError,
18+
SessionAlreadyClosedError,
19+
CursorAlreadyClosedError,
20+
)
1621

1722
logger = logging.getLogger(__name__)
1823

@@ -264,6 +269,12 @@ def _make_request(
264269
error_message = f"Request failed due to retry duration limit: {e}"
265270
# Construct RequestError with message, context, and specific error (like Thrift backend)
266271
raise RequestError(error_message, None, e)
272+
except (SessionAlreadyClosedError, CursorAlreadyClosedError) as e:
273+
# These exceptions are raised by DatabricksRetryPolicy when detecting
274+
# "already closed" scenarios (404 responses with retry history)
275+
error_message = f"Request failed: {e}"
276+
# Construct RequestError with proper 3-argument format (message, context, error) like Thrift backend
277+
raise RequestError(error_message, None, e)
267278

268279
logger.debug(f"Response status: {response.status}")
269280

@@ -282,34 +293,51 @@ def _make_request(
282293
# Handle error responses
283294
error_message = f"SEA HTTP request failed with status {response.status}"
284295

296+
# Try to extract additional error details from response, but don't fail if we can't
297+
error_message = self._try_add_error_details_to_message(response, error_message)
298+
299+
raise RequestError(error_message, None)
300+
301+
def _try_add_error_details_to_message(self, response, error_message: str) -> str:
302+
"""
303+
Try to extract error details from response and add to error message.
304+
This method is defensive and will not raise exceptions if parsing fails.
305+
It handles mock objects and malformed responses gracefully.
306+
"""
285307
try:
286-
if response.data:
308+
# Check if response.data exists and is accessible
309+
if not hasattr(response, "data") or response.data is None:
310+
return error_message
311+
312+
# Try to decode the response data
313+
try:
287314
decoded_data = response.data.decode("utf-8")
288-
# Ensure we have a string before attempting JSON parsing
289-
if isinstance(decoded_data, str):
290-
error_details = json.loads(decoded_data)
291-
if isinstance(error_details, dict) and "message" in error_details:
292-
error_message = f"{error_message}: {error_details['message']}"
293-
logger.error(f"Request failed: {error_details}")
294-
else:
295-
# Handle case where decode returns non-string (e.g., MagicMock in tests)
296-
logger.error(
297-
f"Request failed with non-string response data: {type(decoded_data)}"
298-
)
299-
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
300-
# Log raw response if we can't parse JSON or if we get unexpected types
315+
except (AttributeError, UnicodeDecodeError, TypeError):
316+
# response.data might be a mock object or not bytes
317+
return error_message
318+
319+
# Ensure we have a string before attempting JSON parsing
320+
if not isinstance(decoded_data, str):
321+
return error_message
322+
323+
# Try to parse as JSON
301324
try:
302-
content = (
303-
response.data.decode("utf-8", errors="replace")
304-
if response.data
305-
else ""
325+
error_details = json.loads(decoded_data)
326+
if isinstance(error_details, dict) and "message" in error_details:
327+
enhanced_message = f"{error_message}: {error_details['message']}"
328+
logger.error(f"Request failed: {error_details}")
329+
return enhanced_message
330+
except json.JSONDecodeError:
331+
# Not valid JSON, log what we can
332+
logger.debug(
333+
f"Request failed with non-JSON response: {decoded_data[:200]}"
306334
)
307-
logger.error(f"Request failed with non-JSON response: {content}")
308-
except (AttributeError, TypeError):
309-
# Handle case where response.data itself might be a mock
310-
logger.error(f"Request failed with unparseable response data")
311335

312-
raise RequestError(error_message, None)
336+
except Exception:
337+
# Catch-all for any unexpected issues (e.g., mock objects with unexpected behavior)
338+
logger.debug("Could not parse error response data")
339+
340+
return error_message
313341

314342
def _get_command_type_from_path(self, path: str, method: str) -> CommandType:
315343
"""

0 commit comments

Comments
 (0)