|
1 | 1 | from collections import OrderedDict |
2 | 2 | from decimal import Decimal |
3 | 3 | import itertools |
| 4 | +import logging |
4 | 5 | import unittest |
5 | 6 | import pytest |
6 | 7 | from unittest.mock import patch, MagicMock, Mock |
@@ -320,7 +321,7 @@ def test_binary_protocol_preserves_unbounded_string_and_container_limits( |
320 | 321 | import thrift.protocol.TBinaryProtocol |
321 | 322 |
|
322 | 323 | with patch( |
323 | | - "thrift.protocol.TBinaryProtocol.TBinaryProtocol" |
| 324 | + "thrift.protocol.TBinaryProtocol.TBinaryProtocolAccelerated" |
324 | 325 | ) as mock_protocol_class: |
325 | 326 | ThriftDatabricksClient( |
326 | 327 | "foo", |
@@ -367,6 +368,108 @@ def test_binary_protocol_reads_result_larger_than_thrift_default_limit(self): |
367 | 368 |
|
368 | 369 | self.assertEqual(protocol.readBinary(), oversized_payload) |
369 | 370 |
|
| 371 | + @staticmethod |
| 372 | + def _http_client_with_reply(reply_bytes): |
| 373 | + """A real THttpClient whose connection pool returns ``reply_bytes``.""" |
| 374 | + from databricks.sql.auth.thrift_http_client import THttpClient |
| 375 | + |
| 376 | + http_client = THttpClient( |
| 377 | + auth_provider=Mock(), |
| 378 | + uri_or_host="https://example.com/cliservice", |
| 379 | + ssl_options=SSLOptions(), |
| 380 | + ) |
| 381 | + http_client.setCustomHeaders({"User-Agent": "test"}) |
| 382 | + http_client.open() |
| 383 | + mock_resp = Mock(status=200, reason="OK", headers={}) |
| 384 | + mock_resp.read.return_value = reply_bytes |
| 385 | + http_client._THttpClient__pool = Mock() |
| 386 | + http_client._THttpClient__pool.request.return_value = mock_resp |
| 387 | + return http_client, mock_resp |
| 388 | + |
| 389 | + def test_http_client_buffers_reply_and_decodes_with_accelerated_protocol(self): |
| 390 | + """The reply body is read from urllib3 once, then decoded from memory, |
| 391 | + including by the C-accelerated protocol (the transport is a |
| 392 | + CReadableTransport) for payloads above thrift's ~15.6 MiB default limit.""" |
| 393 | + from thrift.protocol import TBinaryProtocol |
| 394 | + from thrift.transport import TTransport |
| 395 | + from databricks.sql.thrift_api.TCLIService import TCLIService |
| 396 | + from databricks.sql.backend.thrift_backend import ( |
| 397 | + THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, |
| 398 | + THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, |
| 399 | + ) |
| 400 | + |
| 401 | + batch = b"x" * (20 * 1024 * 1024) |
| 402 | + resp = ttypes.TFetchResultsResp( |
| 403 | + status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS), |
| 404 | + hasMoreRows=False, |
| 405 | + results=ttypes.TRowSet( |
| 406 | + startRowOffset=0, |
| 407 | + rows=[], |
| 408 | + arrowBatches=[ttypes.TSparkArrowBatch(batch=batch, rowCount=1)], |
| 409 | + ), |
| 410 | + ) |
| 411 | + reply = TTransport.TMemoryBuffer() |
| 412 | + reply_proto = TBinaryProtocol.TBinaryProtocol(reply) |
| 413 | + reply_proto.writeMessageBegin("FetchResults", 2, 0) # REPLY |
| 414 | + TCLIService.FetchResults_result(success=resp).write(reply_proto) |
| 415 | + reply_proto.writeMessageEnd() |
| 416 | + |
| 417 | + http_client, mock_resp = self._http_client_with_reply(reply.getvalue()) |
| 418 | + self.assertIsInstance(http_client, TTransport.CReadableTransport) |
| 419 | + protocol = TBinaryProtocol.TBinaryProtocolAccelerated( |
| 420 | + http_client, |
| 421 | + string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, |
| 422 | + container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, |
| 423 | + ) |
| 424 | + client = TCLIService.Client(protocol) |
| 425 | + |
| 426 | + op_handle = ttypes.TOperationHandle( |
| 427 | + operationId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16), |
| 428 | + operationType=ttypes.TOperationType.EXECUTE_STATEMENT, |
| 429 | + hasResultSet=True, |
| 430 | + ) |
| 431 | + result = client.FetchResults( |
| 432 | + ttypes.TFetchResultsReq(operationHandle=op_handle, orientation=0, maxRows=1) |
| 433 | + ) |
| 434 | + |
| 435 | + self.assertEqual(result.results.arrowBatches[0].batch, batch) |
| 436 | + mock_resp.read.assert_called_once_with() |
| 437 | + |
| 438 | + http_client.close() |
| 439 | + self.assertEqual(http_client.read(1), b"") |
| 440 | + |
| 441 | + def test_http_client_cstringio_refill_on_truncated_reply_raises(self): |
| 442 | + http_client, _ = self._http_client_with_reply(b"") |
| 443 | + with self.assertRaises(EOFError): |
| 444 | + http_client.cstringio_refill(b"", 4) |
| 445 | + |
| 446 | + @patch("databricks.sql.backend.thrift_backend.unsafe_logger") |
| 447 | + @patch("databricks.sql.auth.thrift_http_client.THttpClient") |
| 448 | + def test_unsafe_logger_formats_payloads_lazily( |
| 449 | + self, t_http_client_class, mock_unsafe_logger |
| 450 | + ): |
| 451 | + """unsafe_logger is DEBUG-level with only a NullHandler by default, so |
| 452 | + payloads must be passed as %s args rather than pre-formatted strings.""" |
| 453 | + thrift_backend = ThriftDatabricksClient( |
| 454 | + "foobar", |
| 455 | + 443, |
| 456 | + "path", |
| 457 | + [], |
| 458 | + auth_provider=AuthProvider(), |
| 459 | + ssl_options=SSLOptions(), |
| 460 | + http_client=MagicMock(), |
| 461 | + ) |
| 462 | + request = Mock() |
| 463 | + response = Mock() |
| 464 | + response.status.statusCode = ttypes.TStatusCode.SUCCESS_STATUS |
| 465 | + method = Mock(return_value=response) |
| 466 | + method.__name__ = "method name" |
| 467 | + |
| 468 | + thrift_backend.make_request(method, request) |
| 469 | + |
| 470 | + mock_unsafe_logger.debug.assert_any_call("Sending request: %s", request) |
| 471 | + mock_unsafe_logger.debug.assert_any_call("Received response: %s", response) |
| 472 | + |
370 | 473 | def test_proxy_headers_are_set(self): |
371 | 474 |
|
372 | 475 | from databricks.sql.common.http_utils import create_basic_proxy_auth_headers |
|
0 commit comments