From e6ffbb68a2d30e5213b1757e98a5d128101b7945 Mon Sep 17 00:00:00 2001 From: anupamme Date: Wed, 5 Aug 2026 01:38:28 +0000 Subject: [PATCH 1/5] fix: V-003 security vulnerability Automated security fix generated by OrbisAI Security --- python/pyspark/sql/connect/client/reattach.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/pyspark/sql/connect/client/reattach.py b/python/pyspark/sql/connect/client/reattach.py index f1d06320866e0..5ed82054c3d2c 100644 --- a/python/pyspark/sql/connect/client/reattach.py +++ b/python/pyspark/sql/connect/client/reattach.py @@ -109,7 +109,9 @@ def __init__( # Initial iterator comes from ExecutePlan request. # Note: This is not retried, because no error would ever be thrown here, and GRPC will only # throw error on first self._has_next(). - self._metadata = metadata + # Convert metadata to a list to ensure it remains re-iterable across all RPCs + # (ReattachExecute, ReleaseExecute), so auth headers are always present. + self._metadata = list(metadata) with disable_gc(): self._iterator: Optional[Iterator[pb2.ExecutePlanResponse]] = iter( self._stub.ExecutePlan( From f130aac0c18a9cf45d92d2d43cfc19bfbe69b61b Mon Sep 17 00:00:00 2001 From: Anupam Mediratta Date: Wed, 5 Aug 2026 14:06:13 +0530 Subject: [PATCH 2/5] [SPARK-57785][CONNECT][PYTHON] Add test for metadata generator exhaustion fix Add a unit test that verifies a single-use generator passed as `metadata` to `ExecutePlanResponseReattachableIterator` is preserved across all subsequent RPCs (`ReattachExecute`, `ReleaseExecute`), not exhausted on the first call. Also extend `MockSparkConnectStub` to record the metadata kwarg received by each RPC method so tests can assert on it. Co-Authored-By: Claude Sonnet 4.6 --- .../sql/tests/connect/client/test_client.py | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/python/pyspark/sql/tests/connect/client/test_client.py b/python/pyspark/sql/tests/connect/client/test_client.py index b5bf76d86df48..4c08aefb6b113 100644 --- a/python/pyspark/sql/tests/connect/client/test_client.py +++ b/python/pyspark/sql/tests/connect/client/test_client.py @@ -113,16 +113,23 @@ def __init__(self, execute_ops=None, attach_ops=None): self.release_calls = 0 self.release_until_calls = 0 self.attach_calls = 0 + # Metadata recorded per call (each entry is the list passed as metadata=) + self.execute_metadata = [] + self.attach_metadata = [] + self.release_metadata = [] def ExecutePlan(self, *args, **kwargs): self.execute_calls += 1 + self.execute_metadata.append(list(kwargs.get("metadata", []))) return self._execute_ops def ReattachExecute(self, *args, **kwargs): self.attach_calls += 1 + self.attach_metadata.append(list(kwargs.get("metadata", []))) return self._attach_ops def ReleaseExecute(self, req: proto.ReleaseExecuteRequest, *args, **kwargs): + self.release_metadata.append(list(kwargs.get("metadata", []))) if req.HasField("release_all"): self.release_calls += 1 elif req.HasField("release_until"): @@ -1186,6 +1193,39 @@ def raise_with_sql_state(): self.assertEqual(err.getErrorClass(), expected_error_class) self.assertEqual(err.getSqlState(), expected_sql_state) + def test_generator_metadata_preserved_across_rpcs(self): + # A single-use generator passed as metadata must not be exhausted before + # ReattachExecute and ReleaseExecute calls; list() in __init__ prevents this. + expected_header = ("x-auth-token", "secret") + + def non_fatal(): + raise TestException("Non Fatal", grpc.StatusCode.UNAVAILABLE) + + stub = self._stub_with( + [self.response, non_fatal], [self.response, self.finished] + ) + metadata_gen = (x for x in [expected_header]) + + ite = ExecutePlanResponseReattachableIterator( + self.request, stub, self.retrying, metadata_gen + ) + for _ in ite: + pass + + def check(): + self.assertEqual(1, stub.attach_calls) + self.assertEqual(1, stub.release_calls) + # ReattachExecute must have received the header. Without the list() + # fix, self._metadata would be the exhausted generator and this fails. + self.assertEqual(1, len(stub.attach_metadata)) + self.assertIn(expected_header, stub.attach_metadata[0]) + # Every ReleaseExecute call must also carry the header. + self.assertGreater(len(stub.release_metadata), 0) + for meta in stub.release_metadata: + self.assertIn(expected_header, meta) + + eventually(timeout=1, catch_assertions=True)(check)() + if __name__ == "__main__": from pyspark.testing import main From 504e7880a722650bcc9e13afb5a7603b7f6107c8 Mon Sep 17 00:00:00 2001 From: Anupam Mediratta Date: Wed, 5 Aug 2026 15:34:20 +0530 Subject: [PATCH 3/5] [SPARK-57785][CONNECT][PYTHON] Fix initial ExecutePlan to use materialized metadata The previous fix materialized metadata into self._metadata via list(metadata) but the initial ExecutePlan call still passed the raw metadata parameter. For a generator input, list(metadata) exhausts it first, so ExecutePlan would receive an empty iterator. Use self._metadata consistently for all RPCs. Update the generator exhaustion test to also assert the initial ExecutePlan received the header. Co-Authored-By: Claude Sonnet 4.6 --- python/pyspark/sql/connect/client/reattach.py | 2 +- python/pyspark/sql/tests/connect/client/test_client.py | 9 ++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/connect/client/reattach.py b/python/pyspark/sql/connect/client/reattach.py index 5ed82054c3d2c..fdf87ceeefabf 100644 --- a/python/pyspark/sql/connect/client/reattach.py +++ b/python/pyspark/sql/connect/client/reattach.py @@ -116,7 +116,7 @@ def __init__( self._iterator: Optional[Iterator[pb2.ExecutePlanResponse]] = iter( self._stub.ExecutePlan( self._initial_request, - metadata=metadata, + metadata=self._metadata, timeout=self._reattachable_execute_plan_timeout, ) ) diff --git a/python/pyspark/sql/tests/connect/client/test_client.py b/python/pyspark/sql/tests/connect/client/test_client.py index 4c08aefb6b113..64216c16e1efe 100644 --- a/python/pyspark/sql/tests/connect/client/test_client.py +++ b/python/pyspark/sql/tests/connect/client/test_client.py @@ -1215,11 +1215,14 @@ def non_fatal(): def check(): self.assertEqual(1, stub.attach_calls) self.assertEqual(1, stub.release_calls) - # ReattachExecute must have received the header. Without the list() - # fix, self._metadata would be the exhausted generator and this fails. + # All three RPC types must receive the header. If list(metadata) on + # line 114 runs before ExecutePlan (which it does), but ExecutePlan + # still uses the raw `metadata` parameter instead of self._metadata, + # then passing a generator would leave execute_metadata[0] empty. + self.assertEqual(1, len(stub.execute_metadata)) + self.assertIn(expected_header, stub.execute_metadata[0]) self.assertEqual(1, len(stub.attach_metadata)) self.assertIn(expected_header, stub.attach_metadata[0]) - # Every ReleaseExecute call must also carry the header. self.assertGreater(len(stub.release_metadata), 0) for meta in stub.release_metadata: self.assertIn(expected_header, meta) From 9e512b61d8157fe2f0894a07384847be5072242f Mon Sep 17 00:00:00 2001 From: Anupam Mediratta Date: Wed, 5 Aug 2026 20:09:12 +0530 Subject: [PATCH 4/5] [SPARK-57785][CONNECT][PYTHON] Fix type hint for self._metadata Add List to the typing imports and annotate self._metadata as List[Tuple[str, str]] to reflect that list(metadata) always produces a list, not a generic Iterable. Co-Authored-By: Claude Sonnet 4.6 --- python/pyspark/sql/connect/client/reattach.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/pyspark/sql/connect/client/reattach.py b/python/pyspark/sql/connect/client/reattach.py index fdf87ceeefabf..475553887a94b 100644 --- a/python/pyspark/sql/connect/client/reattach.py +++ b/python/pyspark/sql/connect/client/reattach.py @@ -19,7 +19,7 @@ from threading import RLock import uuid from collections.abc import Generator -from typing import Optional, Any, Iterator, Iterable, Tuple, Callable, cast, ClassVar +from typing import Optional, Any, Iterator, Iterable, List, Tuple, Callable, cast, ClassVar from concurrent.futures import Future, ThreadPoolExecutor import os import weakref @@ -111,7 +111,7 @@ def __init__( # throw error on first self._has_next(). # Convert metadata to a list to ensure it remains re-iterable across all RPCs # (ReattachExecute, ReleaseExecute), so auth headers are always present. - self._metadata = list(metadata) + self._metadata: List[Tuple[str, str]] = list(metadata) with disable_gc(): self._iterator: Optional[Iterator[pb2.ExecutePlanResponse]] = iter( self._stub.ExecutePlan( From c121dee0071de654c948b2a6271fdcce12fff2ba Mon Sep 17 00:00:00 2001 From: Anupam Mediratta Date: Thu, 6 Aug 2026 06:58:37 +0530 Subject: [PATCH 5/5] [SPARK-57785][CONNECT][PYTHON] Fix all metadata type hints across connect client Tighten metadata type annotations across the three connect client files that form the type chain: - ChannelBuilder.metadata() return type: Iterable -> List - ExecutePlanResponseReattachableIterator.__init__ param: Iterable -> List - ArtifactManager.__init__ param: Iterable -> List - ArtifactManager._metadata attribute: add List annotation + list() copy Co-Authored-By: Claude Sonnet 4.6 --- python/pyspark/sql/connect/client/artifact.py | 4 ++-- python/pyspark/sql/connect/client/core.py | 2 +- python/pyspark/sql/connect/client/reattach.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/connect/client/artifact.py b/python/pyspark/sql/connect/client/artifact.py index 94879171b41f0..2ab1ee88fe7dd 100644 --- a/python/pyspark/sql/connect/client/artifact.py +++ b/python/pyspark/sql/connect/client/artifact.py @@ -168,7 +168,7 @@ def __init__( user_id: Optional[str], session_id: str, channel: grpc.Channel, - metadata: Iterable[Tuple[str, str]], + metadata: List[Tuple[str, str]], add_artifacts_timeout: Optional[float] = None, artifact_status_timeout: Optional[float] = None, ): @@ -177,7 +177,7 @@ def __init__( self._user_context.user_id = user_id self._stub = grpc_lib.SparkConnectServiceStub(channel) self._session_id = session_id - self._metadata = metadata + self._metadata: List[Tuple[str, str]] = list(metadata) self._add_artifacts_timeout = add_artifacts_timeout self._artifact_status_timeout = artifact_status_timeout diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py index 43a22c4998f2b..c3628c4d110e2 100644 --- a/python/pyspark/sql/connect/client/core.py +++ b/python/pyspark/sql/connect/client/core.py @@ -400,7 +400,7 @@ def _effective_channel_options(self) -> List[Tuple[str, Any]]: options.append((key, value)) return options - def metadata(self) -> Iterable[Tuple[str, str]]: + def metadata(self) -> List[Tuple[str, str]]: """ Builds the GRPC specific metadata list to be injected into the request. All parameters will be converted to metadata except ones that are explicitly used diff --git a/python/pyspark/sql/connect/client/reattach.py b/python/pyspark/sql/connect/client/reattach.py index 475553887a94b..0171774ad331b 100644 --- a/python/pyspark/sql/connect/client/reattach.py +++ b/python/pyspark/sql/connect/client/reattach.py @@ -71,7 +71,7 @@ def __init__( request: pb2.ExecutePlanRequest, stub: grpc_lib.SparkConnectServiceStub, retrying: Callable[[], Retrying], - metadata: Iterable[Tuple[str, str]], + metadata: List[Tuple[str, str]], reattachable_execute_plan_timeout: Optional[float] = None, reattach_execute_timeout: Optional[float] = None, ):