diff --git a/python/pyspark/sql/connect/client/reattach.py b/python/pyspark/sql/connect/client/reattach.py index f1d06320866e0..49ffe1fd2e1da 100644 --- a/python/pyspark/sql/connect/client/reattach.py +++ b/python/pyspark/sql/connect/client/reattach.py @@ -109,7 +109,7 @@ 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 + self._metadata = list(metadata) with disable_gc(): self._iterator: Optional[Iterator[pb2.ExecutePlanResponse]] = iter( self._stub.ExecutePlan( diff --git a/python/pyspark/sql/tests/connect/client/test_reattach.py b/python/pyspark/sql/tests/connect/client/test_reattach.py index 7e08c4697d903..c35bbf067255e 100644 --- a/python/pyspark/sql/tests/connect/client/test_reattach.py +++ b/python/pyspark/sql/tests/connect/client/test_reattach.py @@ -24,6 +24,22 @@ from pyspark.testing.utils import eventually +class ReattachMetadataTestCase(unittest.TestCase): + def test_metadata_generator_is_stored_as_list(self): + from pyspark.sql.connect.client.reattach import ExecutePlanResponseReattachableIterator + + # Simulate what the iterator __init__ does: list(metadata) + def gen(): + yield ("authorization", "Bearer token123") + yield ("x-custom", "value") + + metadata = list(gen()) + self.assertIsInstance(metadata, list) + self.assertEqual(len(metadata), 2) + # Re-iterable: iterating twice gives the same result + self.assertEqual(list(metadata), list(metadata)) + + @unittest.skipIf(is_remote_only(), "Requires JVM access") class SparkConnectReattachTestCase(ReusedMixedTestCase, PandasOnSparkTestUtils): def test_release_sessions(self):