diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fbb64b68..d9c58ebf3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ # Unreleased - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) +- Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 0108cc88d..72a20bbf6 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -11,6 +11,7 @@ from http.client import HTTPResponse from io import BytesIO +from thrift.transport.TTransport import CReadableTransport from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager from urllib3.util import make_headers from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy @@ -22,7 +23,7 @@ logger = logging.getLogger(__name__) -class THttpClient(thrift.transport.THttpClient.THttpClient): +class THttpClient(thrift.transport.THttpClient.THttpClient, CReadableTransport): realhost: Optional[str] realport: Optional[int] proxy_uri: Optional[str] @@ -94,6 +95,7 @@ def __init__( self.retry_policy = retry_policy self.__wbuf = BytesIO() + self.__rbuf = BytesIO() self.__resp: Union[None, HTTPResponse] = None self.__timeout = None self.__custom_headers = None @@ -153,9 +155,38 @@ def close(self): self.__resp and self.__resp.drain_conn() self.__resp and self.__resp.release_conn() self.__resp = None + self.__rbuf = BytesIO() + + # Thrift decodes the reply one field at a time, and one urllib3 read per field + # costs far more than the decode, so the body is read in chunks of this size. + _READ_CHUNK_SIZE = 1024 * 1024 def read(self, sz): - return self.__resp.read(sz) + ret = self.__rbuf.read(sz) + if ret: + return ret + self.__rbuf = BytesIO(self.__resp.read(max(sz, self._READ_CHUNK_SIZE))) + return self.__rbuf.read(sz) + + # CReadableTransport: lets TBinaryProtocolAccelerated decode the chunks in C. + @property + def cstringio_buf(self): + return self.__rbuf + + def cstringio_refill(self, partialread, reqlen): + # As in TBufferedTransport: read one more chunk, or all of a larger field. + parts, have = [partialread], len(partialread) + if reqlen < self._READ_CHUNK_SIZE: + parts.append(self.__resp.read(self._READ_CHUNK_SIZE)) + have += len(parts[-1]) + while have < reqlen: + chunk = self.__resp.read(reqlen - have) + if not chunk: + raise EOFError() + parts.append(chunk) + have += len(chunk) + self.__rbuf = BytesIO(b"".join(parts)) + return self.__rbuf def isOpen(self): return self.__resp is not None @@ -206,6 +237,7 @@ def flush(self): self.code = self.__resp.status self.message = self.__resp.reason self.headers = self.__resp.headers + self.__rbuf = BytesIO() logger.info( "HTTP Response with status code {}, message: {}".format( diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 19e378292..7fefd655f 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -253,7 +253,8 @@ def __init__( self._transport.setTimeout(timeout and (float(timeout) * 1000.0)) self._transport.setCustomHeaders(dict(http_headers)) - protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol( + # Uses thrift's C extension when available, else the pure-Python protocol. + protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocolAccelerated( self._transport, string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, @@ -443,7 +444,9 @@ def attempt_request(attempt): this_method_name = getattr(method, "__name__") logger.debug("Sending request: {}()".format(this_method_name)) - unsafe_logger.debug("Sending request: {}".format(request)) + # Lazy %s args: unsafe_logger is DEBUG-level with a NullHandler by + # default, so eager formatting would stringify every payload. + unsafe_logger.debug("Sending request: %s", request) # These three lines are no-ops if the v3 retry policy is not in use if self.enable_v3_retries: @@ -457,7 +460,7 @@ def attempt_request(attempt): logger.debug( "Received response: {}()".format(type(response).__name__) ) - unsafe_logger.debug("Received response: {}".format(response)) + unsafe_logger.debug("Received response: %s", response) return response except urllib3.exceptions.HTTPError as err: diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index dedadaeb2..eb4c26d82 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -1,6 +1,8 @@ from collections import OrderedDict from decimal import Decimal +from io import BytesIO import itertools +import logging import unittest import pytest from unittest.mock import patch, MagicMock, Mock @@ -21,6 +23,13 @@ from databricks.sql.result_set import ResultSet, ThriftResultSet from databricks.sql.backend.types import CommandId, CommandState, SessionId, BackendType +try: + from thrift.protocol import fastbinary # noqa: F401 + + _HAS_THRIFT_FASTBINARY = True +except ImportError: + _HAS_THRIFT_FASTBINARY = False + def retry_policy_factory(): return { # (type, default, min, max) @@ -320,7 +329,7 @@ def test_binary_protocol_preserves_unbounded_string_and_container_limits( import thrift.protocol.TBinaryProtocol with patch( - "thrift.protocol.TBinaryProtocol.TBinaryProtocol" + "thrift.protocol.TBinaryProtocol.TBinaryProtocolAccelerated" ) as mock_protocol_class: ThriftDatabricksClient( "foo", @@ -367,6 +376,341 @@ def test_binary_protocol_reads_result_larger_than_thrift_default_limit(self): self.assertEqual(protocol.readBinary(), oversized_payload) + @staticmethod + def _http_client_with_reply(reply_bytes): + """A real THttpClient whose connection pool streams ``reply_bytes``.""" + from databricks.sql.auth.thrift_http_client import THttpClient + + http_client = THttpClient( + auth_provider=Mock(), + uri_or_host="https://example.com/cliservice", + ssl_options=SSLOptions(), + ) + http_client.setCustomHeaders({"User-Agent": "test"}) + http_client.open() + mock_resp = Mock(status=200, reason="OK", headers={}) + mock_resp.read.side_effect = BytesIO(reply_bytes).read + http_client._THttpClient__pool = Mock() + http_client._THttpClient__pool.request.return_value = mock_resp + return http_client, mock_resp + + @staticmethod + def _fetch_results_reply(batches): + """Wire bytes of a FetchResults REPLY carrying inline Arrow ``batches``.""" + from thrift.protocol import TBinaryProtocol + from thrift.transport import TTransport + from databricks.sql.thrift_api.TCLIService import TCLIService + + resp = ttypes.TFetchResultsResp( + status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS), + hasMoreRows=False, + results=ttypes.TRowSet( + startRowOffset=0, + rows=[], + arrowBatches=[ + ttypes.TSparkArrowBatch(batch=batch, rowCount=1) + for batch in batches + ], + ), + ) + reply = TTransport.TMemoryBuffer() + reply_proto = TBinaryProtocol.TBinaryProtocol(reply) + reply_proto.writeMessageBegin("FetchResults", 2, 0) # REPLY + TCLIService.FetchResults_result(success=resp).write(reply_proto) + reply_proto.writeMessageEnd() + return reply.getvalue() + + @staticmethod + def _fetch_results_req(): + op_handle = ttypes.TOperationHandle( + operationId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16), + operationType=ttypes.TOperationType.EXECUTE_STATEMENT, + hasResultSet=True, + ) + return ttypes.TFetchResultsReq( + operationHandle=op_handle, orientation=0, maxRows=1 + ) + + @staticmethod + def _protocol_client(http_client, protocol_class): + from databricks.sql.thrift_api.TCLIService import TCLIService + from databricks.sql.backend.thrift_backend import ( + THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, + THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, + ) + + protocol = protocol_class( + http_client, + string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, + container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, + ) + return protocol, TCLIService.Client(protocol) + + @staticmethod + def _protocol_classes(): + """The C-accelerated protocol the backend uses, and the pure-Python + protocol it falls back to when thrift's C extension isn't built.""" + from thrift.protocol import TBinaryProtocol + + return ( + TBinaryProtocol.TBinaryProtocolAccelerated, + TBinaryProtocol.TBinaryProtocol, + ) + + def test_http_client_reads_reply_in_chunks(self): + """The reply is pulled from urllib3 in a few chunk-sized reads: not one + read per field, and never the whole body at once. Includes a field + larger than a chunk and than thrift's ~15.6 MiB default limit.""" + from thrift.transport import TTransport + + small = [bytes([i % 256]) * 1000 for i in range(3000)] + big = b"x" * (20 * 1024 * 1024) + reply = self._fetch_results_reply(small + [big]) + for protocol_class in self._protocol_classes(): + with self.subTest(protocol=protocol_class.__name__): + http_client, mock_resp = self._http_client_with_reply(reply) + self.assertIsInstance(http_client, TTransport.CReadableTransport) + _, client = self._protocol_client(http_client, protocol_class) + + result = client.FetchResults(self._fetch_results_req()) + + self.assertEqual( + [b.batch for b in result.results.arrowBatches], small + [big] + ) + read_sizes = [c.args[0] for c in mock_resp.read.call_args_list] + self.assertLess(len(read_sizes), 10) + self.assertLessEqual(max(read_sizes), len(big)) + + @unittest.skipUnless( + _HAS_THRIFT_FASTBINARY, "thrift C extension (fastbinary) not built" + ) + def test_http_client_reply_is_decoded_by_thrift_c_extension(self): + """TBinaryProtocolAccelerated silently falls back to pure Python without + fastbinary, so assert the C decoder actually ran on the streamed body.""" + from thrift.protocol import TBinaryProtocol + + batch = b"x" * (20 * 1024 * 1024) + http_client, _ = self._http_client_with_reply( + self._fetch_results_reply([batch]) + ) + protocol, client = self._protocol_client( + http_client, TBinaryProtocol.TBinaryProtocolAccelerated + ) + protocol._fast_decode = Mock(wraps=protocol._fast_decode) + + result = client.FetchResults(self._fetch_results_req()) + + self.assertEqual(result.results.arrowBatches[0].batch, batch) + protocol._fast_decode.assert_called() + + def test_http_client_decodes_reply_from_short_urllib3_reads(self): + """urllib3 may return fewer bytes than asked for; a field spanning many + short reads must still be assembled correctly.""" + small = [bytes([i % 256]) * 1000 for i in range(300)] + big = b"y" * (4 * 1024 * 1024) + reply = self._fetch_results_reply(small + [big]) + for protocol_class in self._protocol_classes(): + with self.subTest(protocol=protocol_class.__name__): + http_client, mock_resp = self._http_client_with_reply(b"") + stream = BytesIO(reply) + mock_resp.read.side_effect = lambda amt: stream.read(min(amt, 65536)) + _, client = self._protocol_client(http_client, protocol_class) + + result = client.FetchResults(self._fetch_results_req()) + + self.assertEqual( + [b.batch for b in result.results.arrowBatches], small + [big] + ) + + def test_http_client_truncated_reply_raises_eof(self): + reply = self._fetch_results_reply([b"x" * 100]) + for protocol_class in self._protocol_classes(): + with self.subTest(protocol=protocol_class.__name__): + http_client, _ = self._http_client_with_reply(reply[:40]) + _, client = self._protocol_client(http_client, protocol_class) + with self.assertRaises(EOFError): + client.FetchResults(self._fetch_results_req()) + + @patch("time.sleep") + def test_truncated_503_reply_is_retried(self, mock_sleep): + """A 503 whose body is cut short must still reach the status-based retry. + urllib3 1.26 raises on an unbounded read() of a truncated body, but not on + read(amt), so the body must only be read in bounded chunks.""" + import urllib3 + + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [("User-Agent", "test")], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + _enable_v3_retries=False, + ) + truncated_body = BytesIO(b"Service Unavail") + + def urllib3_1_26_read(amt=None): + if amt is None: + raise urllib3.exceptions.ProtocolError("IncompleteRead") + return truncated_body.read(amt) + + unavailable = Mock(status=503, reason="", headers={"Retry-After": "1"}) + unavailable.read.side_effect = urllib3_1_26_read + ok = Mock(status=200, reason="OK", headers={}) + ok.read.side_effect = BytesIO(self._fetch_results_reply([b"x" * 100])).read + pool = thrift_backend._transport._THttpClient__pool = Mock() + pool.request.side_effect = [unavailable, ok] + + result = thrift_backend.make_request( + thrift_backend._client.FetchResults, self._fetch_results_req() + ) + + self.assertEqual(result.results.arrowBatches[0].batch, b"x" * 100) + self.assertEqual(pool.request.call_count, 2) + + def test_reply_buffer_is_released_after_successful_make_request(self): + """make_request closes the transport on success too, so the last chunk + of the reply is not kept on the connection after it has been decoded.""" + batch = b"x" * (1024 * 1024) + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [("User-Agent", "test")], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + mock_resp = Mock(status=200, reason="OK", headers={}) + mock_resp.read.side_effect = BytesIO(self._fetch_results_reply([batch])).read + thrift_backend._transport._THttpClient__pool = Mock() + thrift_backend._transport._THttpClient__pool.request.return_value = mock_resp + + result = thrift_backend.make_request( + thrift_backend._client.FetchResults, self._fetch_results_req() + ) + + self.assertEqual(result.results.arrowBatches[0].batch, batch) + self.assertEqual(thrift_backend._transport.cstringio_buf.getvalue(), b"") + + @unittest.skipUnless( + _HAS_THRIFT_FASTBINARY, "thrift C extension (fastbinary) not built" + ) + @patch("databricks.sql.auth.thrift_http_client.THttpClient") + def test_accelerated_protocol_encodes_backend_requests_identically( + self, t_http_client_class + ): + """TBinaryProtocolAccelerated also moves request encoding to C. Requests + built by the backend must serialize byte-for-byte as with pure Python.""" + import datetime + from thrift.protocol import TBinaryProtocol + from thrift.transport import TTransport + from databricks.sql.parameters import native + + class Captured(Exception): + pass + + captured = [] + + def capture(method, request, *args, **kwargs): + captured.append(request) + raise Captured() + + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + thrift_backend.make_request = capture + session_id = SessionId.from_thrift_handle( + ttypes.TSessionHandle( + sessionId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16) + ) + ) + params = [ + native.IntegerParameter(value=42, name="i"), + native.BigIntegerParameter(value=2**40, name="b"), + native.StringParameter(value="héllo", name="s"), + native.DecimalParameter(value=Decimal("12.345"), name="d"), + native.TimestampParameter( + value=datetime.datetime(2026, 11, 11, 1, 2, 3), name="t" + ), + native.DateParameter(value=datetime.date(2026, 11, 11), name="dt"), + native.BooleanParameter(value=True, name="bo"), + native.VoidParameter(value=None, name="v"), + native.ArrayParameter(value=[1, 2, 3], name="a"), + native.MapParameter(value={"k": "v"}, name="m"), + ] + calls = [ + lambda: thrift_backend.open_session( + {"use_cached_result": False, "ansi_mode": "true"}, "main", "default" + ), + lambda: thrift_backend.execute_command( + operation="SELECT :i, :s", + session_id=session_id, + max_rows=100000, + max_bytes=104857600, + lz4_compression=True, + cursor=Mock(), + use_cloud_fetch=True, + parameters=[p.as_tspark_param(named=True) for p in params], + async_op=False, + enforce_embedded_schema_correctness=False, + row_limit=10, + query_tags={"team": "peco", "flag": None}, + ), + lambda: thrift_backend.get_tables( + session_id, 100000, 104857600, Mock(), "main", "%", "t%", ["TABLE"] + ), + lambda: thrift_backend.get_columns( + session_id, 100000, 104857600, Mock(), "main", "%", "%", "%" + ), + ] + for call in calls: + with self.assertRaises(Captured): + call() + + self.assertEqual(len(captured), len(calls)) + for request in captured: + pure, accelerated = TTransport.TMemoryBuffer(), TTransport.TMemoryBuffer() + request.write(TBinaryProtocol.TBinaryProtocol(pure)) + request.write(TBinaryProtocol.TBinaryProtocolAccelerated(accelerated)) + self.assertEqual( + accelerated.getvalue(), pure.getvalue(), type(request).__name__ + ) + + @patch("databricks.sql.backend.thrift_backend.unsafe_logger") + @patch("databricks.sql.auth.thrift_http_client.THttpClient") + def test_unsafe_logger_formats_payloads_lazily( + self, t_http_client_class, mock_unsafe_logger + ): + """unsafe_logger is DEBUG-level with only a NullHandler by default, so + payloads must be passed as %s args rather than pre-formatted strings.""" + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + request = Mock() + response = Mock() + response.status.statusCode = ttypes.TStatusCode.SUCCESS_STATUS + method = Mock(return_value=response) + method.__name__ = "method name" + + thrift_backend.make_request(method, request) + + mock_unsafe_logger.debug.assert_any_call("Sending request: %s", request) + mock_unsafe_logger.debug.assert_any_call("Received response: %s", response) + def test_proxy_headers_are_set(self): from databricks.sql.common.http_utils import create_basic_proxy_auth_headers