Fix JSONDecodeError in result_batch._load() when fetching large result sets (#2802)

Co-authored-by: Filip Pawłowski <filip.pawlowski@snowflake.com>

Author: sfc-gh-pczajka

DESCRIPTION.md

@@ -16,6 +16,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Fixed a bug where Azure GET commands would incorrectly set the file status to UPLOADED instead of preserving the DOWNLOADED status during metadata retrieval.
   - Renamed the environment variable for skipping config file permission warnings from `SF_SKIP_WARNING_FOR_READ_PERMISSIONS_ON_CONFIG_FILE` to `SF_SKIP_TOKEN_FILE_PERMISSIONS_VERIFICATION`. The old variable is still supported but emits a deprecation warning.
   - Fixed `unsafe_skip_file_permissions_check` flag not being respected when reading `connections.toml`.
+  - Fixed JSONDecodeError in result_batch._load() when fetching large result sets
 
 - v4.3.0(February 12,2026)
   - Ensured proper list conversion - the converter runs to_snowflake on all list elements.

src/snowflake/connector/result_batch.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import gzip
 import json
 import time
 from base64 import b64decode
@@ -52,6 +53,25 @@
 SSE_C_KEY = "x-amz-server-side-encryption-customer-key"
 SSE_C_AES = "AES256"
 
+_GZIP_MAGIC = b"\x1f\x8b"
+
+
+def _ensure_decompressed(response: Response) -> None:
+    """Decompress the response body if HTTP-level gzip decompression was skipped.
+
+    Cloud storage (S3/GCS/Azure) may serve result-set chunks as raw gzip blobs
+    without a ``Content-Encoding: gzip`` header, in which case urllib3's
+    transparent decompression never activates.  Detect this by inspecting the
+    cached *response.content* for the gzip magic number and decompress in-place.
+    """
+    content = response.content
+    if isinstance(content, bytes) and content[:2] == _GZIP_MAGIC:
+        logger.debug(
+            "Response body starts with gzip magic number but was not "
+            "decompressed by the HTTP stack; decompressing explicitly."
+        )
+        response._content = gzip.decompress(content)
+
 
 def _create_nanoarrow_iterator(
     data: bytes,
@@ -418,6 +438,8 @@ def _download(
         self._metrics[DownloadMetrics.download.value] = (
             download_metric.get_timing_millis()
         )
+
+        _ensure_decompressed(response)
         return response
 
     @abc.abstractmethod

test/integ/test_large_result_set.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import logging
+import threading
 from unittest.mock import Mock
 
 import pytest
@@ -221,3 +222,80 @@ def spy_download(self, connection=None, **kwargs):  # type: ignore[no-self-use]
     # Every ResultBatch download reused the same HTTP configuration values
     for cfg in download_cfgs:
         assert cfg == original_cfg
+
+
+# ---------------------------------------------------------------------------
+# Multichunk JSON gzip decompression regression test
+#
+# Verifies that the _ensure_decompressed safety net in _download() correctly
+# handles the case where the HTTP transport does not decompress gzip chunk
+# responses (missing Content-Encoding: gzip header from cloud storage).
+#
+# The standard Snowflake download path *does* decompress correctly because
+# the server sets Content-Encoding: gzip and urllib3 honours it.  The client
+# failure occurs under environmental conditions we cannot reproduce in CI
+# (e.g. a corporate proxy or CDN stripping Content-Encoding).
+#
+# To prove the safety net works we instrument _download to strip the
+# Content-Encoding header from the *real* server response before urllib3
+# consumes it, forcing urllib3 to skip decompression.
+# ---------------------------------------------------------------------------
+
+
+def _strip_content_encoding(response):
+    """Remove Content-Encoding from a urllib3 raw response so decompression is skipped.
+
+    Must be called *before* response.content is accessed (i.e. before the
+    requests library iterates the body), because once the body is consumed
+    the encoding has already been (or not been) applied.
+    """
+    if hasattr(response, "raw") and response.raw is not None:
+        response.raw.headers.pop("Content-Encoding", None)
+        response.raw.headers.pop("content-encoding", None)
+        response.raw._decoder = None
+        response.raw.decode_content = False
+
+
+@pytest.mark.aws
+@pytest.mark.skipolddriver
+def test_multichunk_json_gzip_decompression_with_fix(
+    conn_cnx, db_parameters, ingest_data
+):
+    """With _ensure_decompressed active, multichunk JSON downloads succeed
+    even when the HTTP layer skips gzip decompression."""
+    from snowflake.connector.vendored.requests.adapters import HTTPAdapter
+
+    original_send = HTTPAdapter.send
+    chunks_intercepted = 0
+    chunks_intercepted_lock = threading.Lock()
+
+    def send_strip_encoding(self, request, *args, **kwargs):
+        """Intercept the adapter to strip Content-Encoding after the real
+        HTTP round-trip but before requests reads the body."""
+        nonlocal chunks_intercepted
+        response = original_send(self, request, *args, **kwargs)
+        if response.status_code == 200:
+            _strip_content_encoding(response)
+            with chunks_intercepted_lock:
+                chunks_intercepted += 1
+        return response
+
+    table_name = db_parameters["name"]
+    with conn_cnx(
+        session_parameters={"python_connector_query_result_format": "JSON"},
+    ) as cnx:
+        original_adapter_send = HTTPAdapter.send
+        HTTPAdapter.send = send_strip_encoding
+        try:
+            cur = cnx.cursor()
+            cur.execute(f"select * from {table_name} order by 1")
+            rows = cur.fetchall()
+        finally:
+            HTTPAdapter.send = original_adapter_send
+
+    assert (
+        len(rows) == NUMBER_OF_ROWS
+    ), f"Expected {NUMBER_OF_ROWS} rows, got {len(rows)}"
+    assert (
+        chunks_intercepted > 0
+    ), "No HTTP responses were intercepted -- test didn't exercise the fix"

test/unit/test_result_batch.py

@@ -1,7 +1,10 @@
 #!/usr/bin/env python
 from __future__ import annotations
 
+import gzip
+import json
 from collections import namedtuple
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from http import HTTPStatus
 from test.helpers import create_mock_response
 from unittest import mock
@@ -40,15 +43,22 @@
 try:
     from snowflake.connector.compat import TOO_MANY_REQUESTS
     from snowflake.connector.errors import TooManyRequests
-    from snowflake.connector.result_batch import MAX_DOWNLOAD_RETRY, JSONResultBatch
-    from snowflake.connector.vendored import requests  # NOQA
+    from snowflake.connector.result_batch import (
+        MAX_DOWNLOAD_RETRY,
+        JSONResultBatch,
+        RemoteChunkInfo,
+        _ensure_decompressed,
+    )
+    from snowflake.connector.vendored import requests
 
     SESSION_FROM_REQUEST_MODULE_PATH = (
         "snowflake.connector.vendored.requests.sessions.Session"
     )
 except ImportError:
     MAX_DOWNLOAD_RETRY = None
     JSONResultBatch = None
+    RemoteChunkInfo = None
+    _ensure_decompressed = None
     SESSION_FROM_REQUEST_MODULE_PATH = "requests.sessions.Session"
     TooManyRequests = None
     TOO_MANY_REQUESTS = None
@@ -151,3 +161,184 @@ def test_retries_until_success():
             assert res.raw == "success"
         # call `get` once for each error and one last time when it succeeds
         assert mock_get.call_count == len(error_codes) + 1
+
+
+# ---------------------------------------------------------------------------
+# Gzip decompression fallback tests
+#
+# These reproduce the JSONDecodeError observed when cloud storage serves
+# result-set chunks as raw gzip blobs *without* a Content-Encoding: gzip
+# header.  urllib3 v2 only triggers transparent decompression when that
+# header is present, so the raw \x1f\x8b bytes leak into response.text.
+# ---------------------------------------------------------------------------
+
+
+def _make_gzip_json_rows(*rows):
+    """Encode rows as Snowflake-style comma-separated JSON and gzip-compress."""
+    payload = ",\n".join(json.dumps(row) for row in rows)
+    return gzip.compress(payload.encode("utf-8"))
+
+
+def _make_gzip_response(compressed_body: bytes):
+    """Build a fake requests.Response whose .content is raw gzip bytes.
+
+    This simulates what happens when cloud storage returns gzip data
+    without setting Content-Encoding: gzip -- the requests/urllib3 stack
+    skips decompression and .content returns the raw compressed bytes.
+    """
+    resp = requests.Response()
+    resp.status_code = 200
+    resp._content = compressed_body
+    resp.headers["Content-Type"] = "application/json"
+    return resp
+
+
+@pytest.mark.skipif(JSONResultBatch is None, reason="vendored requests unavailable")
+class TestGzipDecompressionFallback:
+    """Verify _ensure_decompressed fixes responses that were not decoded by urllib3."""
+
+    def test_ensure_decompressed_unpacks_gzip_content(self):
+        """_ensure_decompressed should replace raw gzip bytes with decompressed content."""
+        rows = [["Alice", 30], ["Bob", 25]]
+        raw_gz = _make_gzip_json_rows(*rows)
+        assert raw_gz[:2] == b"\x1f\x8b", "sanity: payload is gzip"
+
+        resp = _make_gzip_response(raw_gz)
+        assert resp.content[:2] == b"\x1f\x8b", "before fix: content is raw gzip"
+
+        _ensure_decompressed(resp)
+
+        assert resp.content[:2] != b"\x1f\x8b", "after fix: gzip magic gone"
+        recovered = json.loads("[" + resp.content.decode("utf-8") + "]")
+        assert recovered == rows
+
+    def test_ensure_decompressed_leaves_plain_json_alone(self):
+        """_ensure_decompressed should be a no-op for already-decoded responses."""
+        plain = b'["Alice", 30],\n["Bob", 25]'
+        resp = requests.Response()
+        resp.status_code = 200
+        resp._content = plain
+        original_id = id(resp._content)
+
+        _ensure_decompressed(resp)
+
+        assert resp.content is plain or resp.content == plain
+        assert id(resp._content) == original_id
+
+    def test_json_result_batch_load_with_gzip_response(self):
+        """JSONResultBatch._load should succeed even when the HTTP layer didn't decompress."""
+        rows = [["val1", 1], ["val2", 2], ["val3", 3]]
+        raw_gz = _make_gzip_json_rows(*rows)
+
+        resp = _make_gzip_response(raw_gz)
+        _ensure_decompressed(resp)
+
+        batch = JSONResultBatch(
+            rowcount=len(rows),
+            chunk_headers=None,
+            remote_chunk_info=None,
+            schema=[],
+            column_converters=[],
+            use_dict_result=False,
+        )
+        loaded = batch._load(resp)
+        assert loaded == rows
+
+    def test_concurrent_multichunk_download_with_gzip_responses(self):
+        """Reproduce the reported issue: concurrent ThreadPoolExecutor downloads
+        where each chunk response is raw gzip (no Content-Encoding header).
+
+        Without the _ensure_decompressed fix, json.loads() in _load() would
+        receive \\x1f\\x8b... garbage and raise JSONDecodeError.
+        """
+        num_chunks = 6
+        rows_per_chunk = 50
+        chunks_data = {}
+        url_to_response = {}
+
+        for chunk_idx in range(num_chunks):
+            rows = [
+                [f"chunk{chunk_idx}_row{r}", chunk_idx * 100 + r]
+                for r in range(rows_per_chunk)
+            ]
+            chunk_url = f"http://fake-s3.example.com/results/chunk_{chunk_idx}"
+            body = _make_gzip_json_rows(*rows)
+            chunks_data[chunk_idx] = rows
+            url_to_response[chunk_url] = _make_gzip_response(body)
+
+        batches = []
+        for chunk_idx in range(num_chunks):
+            chunk_url = f"http://fake-s3.example.com/results/chunk_{chunk_idx}"
+            body = url_to_response[chunk_url].content
+            batch = JSONResultBatch(
+                rowcount=rows_per_chunk,
+                chunk_headers=None,
+                remote_chunk_info=RemoteChunkInfo(
+                    url=chunk_url, uncompressedSize=0, compressedSize=len(body)
+                ),
+                schema=[],
+                column_converters=[],
+                use_dict_result=False,
+            )
+            batches.append((chunk_idx, batch))
+
+        def mock_get(url, **kwargs):
+            return url_to_response[url]
+
+        def fetch_batch(idx_and_batch):
+            idx, batch = idx_and_batch
+            response = batch._download()
+            return idx, batch._load(response)
+
+        all_results = {}
+        with mock.patch(
+            SESSION_FROM_REQUEST_MODULE_PATH + ".get", side_effect=mock_get
+        ):
+            with ThreadPoolExecutor(max_workers=4) as pool:
+                futures = [pool.submit(fetch_batch, ib) for ib in batches]
+                for future in as_completed(futures):
+                    chunk_idx, loaded = future.result()
+                    all_results[chunk_idx] = loaded
+
+        assert len(all_results) == num_chunks
+        for chunk_idx in range(num_chunks):
+            assert all_results[chunk_idx] == chunks_data[chunk_idx], (
+                f"Chunk {chunk_idx}: expected valid JSON rows but got corrupted data. "
+                f"This indicates gzip decompression was not applied."
+            )
+
+    def test_concurrent_multichunk_with_session_manager_clone(self):
+        """End-to-end reproduction using a cloned SessionManager, mirroring the
+        real download path where result batches use a cloned manager without
+        connection pooling.
+        """
+        from snowflake.connector.session_manager import SessionManager
+
+        base_manager = SessionManager()
+        cloned_manager = base_manager.clone(use_pooling=False)
+
+        rows = [["hello", 42], ["world", 99]]
+        raw_gz = _make_gzip_json_rows(*rows)
+        chunk_url = "http://fake-s3.example.com/results/chunk_0"
+
+        batch = JSONResultBatch(
+            rowcount=len(rows),
+            chunk_headers={},
+            remote_chunk_info=RemoteChunkInfo(
+                url=chunk_url, uncompressedSize=0, compressedSize=len(raw_gz)
+            ),
+            schema=[],
+            column_converters=[],
+            use_dict_result=False,
+            session_manager=cloned_manager,
+        )
+
+        resp = _make_gzip_response(raw_gz)
+
+        with mock.patch(
+            SESSION_FROM_REQUEST_MODULE_PATH + ".request", return_value=resp
+        ):
+            response = batch._download()
+
+        loaded = batch._load(response)
+        assert loaded == rows