python/pyspark/sql/tests/connect/client/test_client.py - master - githubmirror/spark

python/pyspark/sql/tests/connect/client/test_client.py
1 238 строк49 KB
Wenchen Fan
[SPARK-58547][CONNECT] Expose operation IDs for end-to-end request attribution
05 авг 2026, 12:14
05 авг 2026, 12:14e4043f3
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import unittest
import uuid
from collections.abc import Generator
from typing import Optional, Any, Union

from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
from pyspark.testing.utils import eventually

if should_test_connect:
    import grpc
    import google.protobuf.any_pb2 as any_pb2
    import google.protobuf.wrappers_pb2 as wrappers_pb2
    from google.rpc import status_pb2
    from google.rpc.error_details_pb2 import ErrorInfo
    import pandas as pd
    import pyarrow as pa
    from pyspark.sql.connect.client import SparkConnectClient, DefaultChannelBuilder
    from pyspark.sql.connect.client.core import RpcDeadlines
    from pyspark.sql.connect.client.retries import (
        Retrying,
        DefaultPolicy,
    )
    from pyspark.sql.connect.client.reattach import ExecutePlanResponseReattachableIterator
    from pyspark.sql.connect.session import SparkSession as RemoteSparkSession
    from pyspark.errors import PySparkRuntimeError
    from pyspark.errors.exceptions.connect import SparkConnectGrpcException
    import pyspark.sql.connect.proto as proto

    class TestPolicy(DefaultPolicy):
        def __init__(self):
            super().__init__(
                max_retries=3,
                backoff_multiplier=4.0,
                initial_backoff=10,
                max_backoff=10,
                jitter=10,
                min_jitter_threshold=10,
            )

    class TestException(grpc.RpcError, grpc.Call):
        """Exception mock to test retryable exceptions."""

        def __init__(
            self,
            msg,
            code=grpc.StatusCode.INTERNAL,
            trailing_status: Union[status_pb2.Status, None] = None,
        ):
            self.msg = msg
            self._code = code
            self._trailer: dict[str, Any] = {}
            if trailing_status is not None:
                self._trailer["grpc-status-details-bin"] = trailing_status.SerializeToString()

        def code(self):
            return self._code

        def __str__(self):
            return self.msg

        def details(self):
            return self.msg

        def trailing_metadata(self):
            return None if not self._trailer else self._trailer.items()

    class ResponseGenerator(Generator):
        """This class is used to generate values that are returned by the streaming
        iterator of the GRPC stub."""

        def __init__(self, funs):
            self._funs = funs
            self._iterator = iter(self._funs)

        def send(self, value: Any) -> proto.ExecutePlanResponse:
            val = next(self._iterator)
            if callable(val):
                return val()
            else:
                return val

        def throw(self, type: Any = None, value: Any = None, traceback: Any = None) -> Any:
            super().throw(type, value, traceback)

        def close(self) -> None:
            return super().close()

    class MockSparkConnectStub:
        """Simple mock class for the GRPC stub used by the re-attachable execution."""

        def __init__(self, execute_ops=None, attach_ops=None):
            self._execute_ops = execute_ops
            self._attach_ops = attach_ops
            # Call counters
            self.execute_calls = 0
            self.release_calls = 0
            self.release_until_calls = 0
            self.attach_calls = 0

        def ExecutePlan(self, *args, **kwargs):
            self.execute_calls += 1
            return self._execute_ops

        def ReattachExecute(self, *args, **kwargs):
            self.attach_calls += 1
            return self._attach_ops

        def ReleaseExecute(self, req: proto.ReleaseExecuteRequest, *args, **kwargs):
            if req.HasField("release_all"):
                self.release_calls += 1
            elif req.HasField("release_until"):
                self.release_until_calls += 1

    class MockService:
        # Simplest mock of the SparkConnectService.
        # If this needs more complex logic, it needs to be replaced with Python mocking.

        req: Optional[proto.ExecutePlanRequest]

        OperationStatus = proto.GetStatusResponse.OperationStatus
        DEFAULT_OPERATION_STATUSES = [
            OperationStatus(
                operation_id="default-op-1",
                state=OperationStatus.OperationState.OPERATION_STATE_SUCCEEDED,
            ),
            OperationStatus(
                operation_id="default-op-2",
                state=OperationStatus.OperationState.OPERATION_STATE_RUNNING,
            ),
        ]

        def __init__(self, session_id: str, operation_statuses=None):
            self._session_id = session_id
            self.req = None
            self.client_user_context_extensions = []
            if operation_statuses is None:
                operation_statuses = self.DEFAULT_OPERATION_STATUSES
            self._operation_statuses = {s.operation_id: s for s in operation_statuses}

        def ExecutePlan(self, req: proto.ExecutePlanRequest, metadata, timeout=None):
            self.req = req
            self.client_user_context_extensions = list(req.user_context.extensions)
            resp = proto.ExecutePlanResponse()
            resp.session_id = self._session_id
            resp.operation_id = req.operation_id

            pdf = pd.DataFrame(data={"col1": [1, 2]})
            schema = pa.Schema.from_pandas(pdf)
            table = pa.Table.from_pandas(pdf)
            sink = pa.BufferOutputStream()

            writer = pa.ipc.new_stream(sink, schema=schema)
            writer.write(table)
            writer.close()

            buf = sink.getvalue()
            resp.arrow_batch.data = buf.to_pybytes()
            resp.arrow_batch.row_count = 2
            return [resp]

        def Interrupt(self, req: proto.InterruptRequest, metadata, timeout=None):
            self.req = req
            self.client_user_context_extensions = list(req.user_context.extensions)
            resp = proto.InterruptResponse()
            resp.session_id = self._session_id
            return resp

        def Config(self, req: proto.ConfigRequest, metadata, timeout=None):
            self.req = req
            self.client_user_context_extensions = list(req.user_context.extensions)
            resp = proto.ConfigResponse()
            resp.session_id = self._session_id
            if req.operation.HasField("get"):
                pair = resp.pairs.add()
                pair.key = req.operation.get.keys[0]
                pair.value = "true"  # Default value
            elif req.operation.HasField("get_with_default"):
                pair = resp.pairs.add()
                pair.key = req.operation.get_with_default.pairs[0].key
                pair.value = req.operation.get_with_default.pairs[0].value or "true"
            return resp

        def AnalyzePlan(self, req: proto.AnalyzePlanRequest, metadata, timeout=None):
            self.req = req
            self.client_user_context_extensions = list(req.user_context.extensions)
            resp = proto.AnalyzePlanResponse()
            resp.session_id = self._session_id
            # Return a minimal response with a semantic hash
            resp.semantic_hash.result = 12345
            return resp

        def GetStatus(self, req: proto.GetStatusRequest, metadata, timeout=None):
            self.req = req
            self.client_user_context_extensions = list(req.user_context.extensions)
            self.received_custom_server_session_id = req.client_observed_server_side_session_id
            resp = proto.GetStatusResponse(session_id=self._session_id)

            # Echo top-level request extensions back in the response
            if req.extensions:
                resp.extensions.extend(req.extensions)

            if not req.HasField("operation_status"):
                return resp

            # Collect operation-status-level extensions from the request to echo back
            op_status_extensions = list(req.operation_status.extensions)

            requested_ids = list(req.operation_status.operation_ids)
            if len(requested_ids) == 0:
                # Empty list — return all statuses
                resp.operation_statuses.extend(self._operation_statuses.values())
                return resp

            OperationStatus = proto.GetStatusResponse.OperationStatus
            for op_id in requested_ids:
                status = self._operation_statuses.get(op_id)
                if status is not None:
                    op_status = OperationStatus(
                        operation_id=status.operation_id,
                        state=status.state,
                    )
                else:
                    op_status = OperationStatus(
                        operation_id=op_id,
                        state=OperationStatus.OperationState.OPERATION_STATE_UNKNOWN,
                    )
                op_status.extensions.extend(op_status_extensions)
                resp.operation_statuses.append(op_status)
            return resp

    # The _cleanup_ml_cache invocation will hang in this test (no valid spark cluster)
    # and it blocks the test process exiting because it is registered as the atexit handler
    # in `SparkConnectClient` constructor. To bypass the issue, patch the method in the test.
    SparkConnectClient._cleanup_ml_cache = lambda _: None


@unittest.skipIf(not should_test_connect, connect_requirement_message)
class SparkConnectClientTestCase(unittest.TestCase):
    def test_user_agent_passthrough(self):
        client = SparkConnectClient("sc://foo/;user_agent=bar", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock

        command = proto.Command()
        client.execute_command(command)

        self.assertIsNotNone(mock.req, "ExecutePlan API was not called when expected")
        self.assertRegex(mock.req.client_type, r"^bar spark/[^ ]+ os/[^ ]+ python/[^ ]+$")

    def test_user_agent_default(self):
        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock

        command = proto.Command()
        client.execute_command(command)

        self.assertIsNotNone(mock.req, "ExecutePlan API was not called when expected")
        self.assertRegex(
            mock.req.client_type, r"^_SPARK_CONNECT_PYTHON spark/[^ ]+ os/[^ ]+ python/[^ ]+$"
        )

    def test_properties(self):
        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        self.assertEqual(client.token, "bar")
        self.assertEqual(client.host, "foo")

        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        self.assertIsNone(client.token)

    def test_channel_builder(self):
        class CustomChannelBuilder(DefaultChannelBuilder):
            @property
            def userId(self) -> Optional[str]:
                return "abc"

        client = SparkConnectClient(
            CustomChannelBuilder("sc://foo/"), use_reattachable_execute=False
        )

        self.assertEqual(client._user_id, "abc")

    def test_user_context_extension(self):
        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock

        try:
            exlocal = any_pb2.Any()
            exlocal.Pack(wrappers_pb2.StringValue(value="abc"))
            exlocal2 = any_pb2.Any()
            exlocal2.Pack(wrappers_pb2.StringValue(value="def"))
            exglobal = any_pb2.Any()
            exglobal.Pack(wrappers_pb2.StringValue(value="ghi"))
            exglobal2 = any_pb2.Any()
            exglobal2.Pack(wrappers_pb2.StringValue(value="jkl"))

            exlocal_id = client.add_threadlocal_user_context_extension(exlocal)
            exglobal_id = client.add_global_user_context_extension(exglobal)

            mock.client_user_context_extensions = []
            command = proto.Command()
            client.execute_command(command)
            self.assertTrue(exlocal in mock.client_user_context_extensions)
            self.assertTrue(exglobal in mock.client_user_context_extensions)
            self.assertFalse(exlocal2 in mock.client_user_context_extensions)
            self.assertFalse(exglobal2 in mock.client_user_context_extensions)

            client.add_threadlocal_user_context_extension(exlocal2)

            mock.client_user_context_extensions = []
            plan = proto.Plan()
            client.semantic_hash(plan)  # use semantic_hash to test analyze
            self.assertTrue(exlocal in mock.client_user_context_extensions)
            self.assertTrue(exglobal in mock.client_user_context_extensions)
            self.assertTrue(exlocal2 in mock.client_user_context_extensions)
            self.assertFalse(exglobal2 in mock.client_user_context_extensions)

            client.add_global_user_context_extension(exglobal2)

            mock.client_user_context_extensions = []
            client.interrupt_all()
            self.assertTrue(exlocal in mock.client_user_context_extensions)
            self.assertTrue(exglobal in mock.client_user_context_extensions)
            self.assertTrue(exlocal2 in mock.client_user_context_extensions)
            self.assertTrue(exglobal2 in mock.client_user_context_extensions)

            client.remove_user_context_extension(exlocal_id)

            mock.client_user_context_extensions = []
            client.get_configs("foo", "bar")
            self.assertFalse(exlocal in mock.client_user_context_extensions)
            self.assertTrue(exglobal in mock.client_user_context_extensions)
            self.assertTrue(exlocal2 in mock.client_user_context_extensions)
            self.assertTrue(exglobal2 in mock.client_user_context_extensions)

            client.remove_user_context_extension(exglobal_id)

            mock.client_user_context_extensions = []
            command = proto.Command()
            client.execute_command(command)
            self.assertFalse(exlocal in mock.client_user_context_extensions)
            self.assertFalse(exglobal in mock.client_user_context_extensions)
            self.assertTrue(exlocal2 in mock.client_user_context_extensions)
            self.assertTrue(exglobal2 in mock.client_user_context_extensions)

            client.clear_user_context_extensions()

            mock.client_user_context_extensions = []
            plan = proto.Plan()
            client.semantic_hash(plan)  # use semantic_hash to test analyze
            self.assertFalse(exlocal in mock.client_user_context_extensions)
            self.assertFalse(exglobal in mock.client_user_context_extensions)
            self.assertFalse(exlocal2 in mock.client_user_context_extensions)
            self.assertFalse(exglobal2 in mock.client_user_context_extensions)

            mock.client_user_context_extensions = []
            client.interrupt_all()
            self.assertFalse(exlocal in mock.client_user_context_extensions)
            self.assertFalse(exglobal in mock.client_user_context_extensions)
            self.assertFalse(exlocal2 in mock.client_user_context_extensions)
            self.assertFalse(exglobal2 in mock.client_user_context_extensions)

            mock.client_user_context_extensions = []
            client.get_configs("foo", "bar")
            self.assertFalse(exlocal in mock.client_user_context_extensions)
            self.assertFalse(exglobal in mock.client_user_context_extensions)
            self.assertFalse(exlocal2 in mock.client_user_context_extensions)
            self.assertFalse(exglobal2 in mock.client_user_context_extensions)
        finally:
            client.close()

    def test_interrupt_all(self):
        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock

        client.interrupt_all()
        self.assertIsNotNone(mock.req, "Interrupt API was not called when expected")

    def test_is_closed(self):
        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)

        self.assertFalse(client.is_closed)
        client.close()
        self.assertTrue(client.is_closed)

    def test_channel_builder_with_session(self):
        dummy = str(uuid.uuid4())
        chan = DefaultChannelBuilder(f"sc://foo/;session_id={dummy}")
        client = SparkConnectClient(chan)
        self.assertEqual(client._session_id, chan.session_id)

    def test_session_hook(self):
        inits = 0
        calls = 0

        class TestHook(RemoteSparkSession.Hook):
            def __init__(self, _session):
                nonlocal inits
                inits += 1

            def on_execute_plan(self, req):
                nonlocal calls
                calls += 1
                return req

        session = (
            RemoteSparkSession.builder.remote("sc://foo")._registerHook(TestHook).getOrCreate()
        )
        try:
            self.assertEqual(inits, 1)
            self.assertEqual(calls, 0)
            session.client._stub = MockService(session.client._session_id)
            session.client.disable_reattachable_execute()

            # Called from _execute_and_fetch_as_iterator
            session.range(1).collect()
            self.assertEqual(inits, 1)
            self.assertEqual(calls, 1)

            # Called from _execute
            session.udf.register("test_func", lambda x: x + 1)
            self.assertEqual(inits, 1)
            self.assertEqual(calls, 2)
        finally:
            # Close the session to avoid leaking dummy session inter-test
            session.stop()

    def test_session_hook_preserved_after_new_session(self):
        calls = 0

        class TestHook(RemoteSparkSession.Hook):
            def __init__(self, _session):
                pass

            def on_execute_plan(self, req):
                nonlocal calls
                calls += 1
                return req

        # Use create() instead of getOrCreate() to avoid picking up a session (and hooks)
        # left active by other tests.
        session = RemoteSparkSession.builder.remote("sc://foo")._registerHook(TestHook).create()
        new_session = session.newSession()
        try:
            # Client-side behavior carries over to the fresh session, as in clone().
            self.assertEqual(new_session.client._session_hooks, session.client._session_hooks)
            self.assertEqual(new_session.client._rpc_deadlines, session.client._rpc_deadlines)

            new_session.client._stub = MockService(new_session.client._session_id)
            new_session.client.disable_reattachable_execute()

            # The hook still observes ExecutePlanRequests issued through the new session.
            self.assertEqual(calls, 0)
            new_session.range(1).collect()
            self.assertEqual(calls, 1)
        finally:
            # Close the clients so their atexit hooks do not try to release sessions
            # against the unreachable endpoint at interpreter shutdown.
            new_session.client.close()
            session.client.close()
            session.stop()

    def test_session_hook_preserves_operation_id(self):
        execute_plan_req = None

        class TestHook(RemoteSparkSession.Hook):
            def __init__(self, _session):
                pass

            def on_execute_plan(self, req):
                replacement = proto.ExecutePlanRequest()
                replacement.CopyFrom(req)
                replacement.ClearField("operation_id")
                return replacement

        class TestService(MockService):
            def ExecutePlan(self, req, metadata, timeout=None):
                nonlocal execute_plan_req
                execute_plan_req = req
                return super().ExecutePlan(req, metadata, timeout)

        session = (
            RemoteSparkSession.builder.remote("sc://foo")._registerHook(TestHook).getOrCreate()
        )
        try:
            mock = TestService(session.client._session_id)
            session.client._stub = mock
            session.client.disable_reattachable_execute()

            df = session.range(1)
            df.collect()
            self.assertIsNotNone(df.executionInfo)
            self.assertIsNotNone(execute_plan_req)
            assert execute_plan_req is not None
            self.assertEqual(execute_plan_req.operation_id, df.executionInfo.operation_id)
            uuid.UUID(execute_plan_req.operation_id)
        finally:
            session.stop()

    def test_new_session_preserves_custom_channel_builder(self):
        class CustomChannelBuilder(DefaultChannelBuilder):
            pass

        client = SparkConnectClient(
            CustomChannelBuilder("sc://foo/"), use_reattachable_execute=False
        )
        new_client = client.newSession()
        try:
            # newSession() deep-copies the connection configuration, so the builder
            # subclass is preserved and the original builder is left untouched.
            self.assertIsInstance(new_client._builder, CustomChannelBuilder)
            self.assertIsNot(new_client._builder, client._builder)
            # The session id is dropped from the copied parameters and regenerated.
            self.assertIsNone(new_client._builder.session_id)
            self.assertNotEqual(new_client._session_id, client._session_id)
        finally:
            new_client.close()
            client.close()

    def test_custom_operation_id(self):
        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock
        req = client._execute_plan_request_with_metadata(
            operation_id="10a4c38e-7e87-40ee-9d6f-60ff0751e63b"
        )
        for resp in client._stub.ExecutePlan(req, metadata=None):
            assert resp.operation_id == "10a4c38e-7e87-40ee-9d6f-60ff0751e63b"

    def test_execute_plan_request_generates_operation_id(self):
        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        try:
            req = client._execute_plan_request_with_metadata()
            uuid.UUID(req.operation_id)
        finally:
            client.close()

    def test_on_exit_calls_release_and_close_when_enabled(self):
        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        client._release_session_on_exit = True
        client._closed = False

        call_tracker = {"release_session": 0, "close": 0}

        def mock_release_session():
            call_tracker["release_session"] += 1

        def mock_close():
            call_tracker["close"] += 1

        client.release_session = mock_release_session
        client.close = mock_close

        client._on_exit()

        self.assertEqual(call_tracker["release_session"], 1)
        self.assertEqual(call_tracker["close"], 1)

    def test_on_exit_does_not_call_when_release_disabled(self):
        """Test _on_exit does nothing when _release_session_on_exit is False."""
        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        client._release_session_on_exit = False
        client._closed = False

        call_tracker = {"release_session": 0, "close": 0}

        def mock_release_session():
            call_tracker["release_session"] += 1

        def mock_close():
            call_tracker["close"] += 1

        client.release_session = mock_release_session
        client.close = mock_close

        client._on_exit()

        self.assertEqual(call_tracker["release_session"], 0)
        self.assertEqual(call_tracker["close"], 0)

    def test_on_exit_does_not_call_when_already_closed(self):
        """Test _on_exit does nothing when client is already closed."""
        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        client._release_session_on_exit = True
        client._closed = True

        call_tracker = {"cleanup_ml_cache": 0, "release_session": 0, "close": 0}

        def mock_cleanup_ml_cache():
            call_tracker["cleanup_ml_cache"] += 1

        def mock_release_session():
            call_tracker["release_session"] += 1

        def mock_close():
            call_tracker["close"] += 1

        client._cleanup_ml_cache = mock_cleanup_ml_cache
        client.release_session = mock_release_session
        client.close = mock_close

        client._on_exit()

        self.assertEqual(call_tracker["cleanup_ml_cache"], 0)
        self.assertEqual(call_tracker["release_session"], 0)
        self.assertEqual(call_tracker["close"], 0)

    def test_on_exit_catches_release_session_exception(self):
        """Test _on_exit continues to call close even if release_session raises."""
        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        client._release_session_on_exit = True
        client._closed = False

        call_tracker = {"release_session": 0, "close": 0}

        def mock_release_session():
            call_tracker["release_session"] += 1
            raise Exception("release error")

        def mock_close():
            call_tracker["close"] += 1

        client.release_session = mock_release_session
        client.close = mock_close

        # Should not raise
        client._on_exit()

        self.assertEqual(call_tracker["release_session"], 1)
        self.assertEqual(call_tracker["close"], 1)

    def test_on_exit_catches_close_exception(self):
        """Test _on_exit silently catches exception from close."""
        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        client._release_session_on_exit = True
        client._closed = False

        call_tracker = {"release_session": 0, "close": 0}

        def mock_release_session():
            call_tracker["release_session"] += 1

        def mock_close():
            call_tracker["close"] += 1
            raise Exception("close error")

        client.release_session = mock_release_session
        client.close = mock_close

        # Should not raise
        client._on_exit()

        self.assertEqual(call_tracker["release_session"], 1)
        self.assertEqual(call_tracker["close"], 1)

    def test_on_exit_catches_both_exceptions(self):
        """Test _on_exit handles both release_session and close raising exceptions."""
        client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
        client._release_session_on_exit = True
        client._closed = False

        call_tracker = {"release_session": 0, "close": 0}

        def mock_release_session():
            call_tracker["release_session"] += 1
            raise Exception("release error")

        def mock_close():
            call_tracker["close"] += 1
            raise Exception("close error")

        client.release_session = mock_release_session
        client.close = mock_close

        # Should not raise
        client._on_exit()

        self.assertEqual(call_tracker["release_session"], 1)
        self.assertEqual(call_tracker["close"], 1)

    def test_get_operations_statuses_all(self):
        """Test get_operations_statuses returns all operation statuses when no IDs specified."""
        OperationStatus = proto.GetStatusResponse.OperationStatus
        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock

        resp = client._get_operation_statuses()
        result = list(resp.operation_statuses)
        self.assertEqual(len(result), 2)
        status_map = {s.operation_id: s.state for s in result}
        self.assertEqual(
            status_map["default-op-1"],
            OperationStatus.OperationState.OPERATION_STATE_SUCCEEDED,
        )
        self.assertEqual(
            status_map["default-op-2"],
            OperationStatus.OperationState.OPERATION_STATE_RUNNING,
        )

    def test_get_operations_statuses_specific_ids(self):
        """Test get_operations_statuses filters by specific operation IDs."""
        OperationStatus = proto.GetStatusResponse.OperationStatus
        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock

        resp = client._get_operation_statuses(operation_ids=["default-op-1", "unknown-op"])
        result = list(resp.operation_statuses)
        self.assertEqual(len(result), 2)
        status_map = {s.operation_id: s.state for s in result}
        self.assertEqual(
            status_map["default-op-1"],
            OperationStatus.OperationState.OPERATION_STATE_SUCCEEDED,
        )
        self.assertEqual(
            status_map["unknown-op"],
            OperationStatus.OperationState.OPERATION_STATE_UNKNOWN,
        )
        # Verify the request included the operation IDs
        self.assertEqual(
            set(mock.req.operation_status.operation_ids), {"default-op-1", "unknown-op"}
        )

    def test_get_operations_statuses_empty(self):
        """Test get_operations_statuses returns empty list when no operations exist."""
        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        mock = MockService(client._session_id, operation_statuses=[])
        client._stub = mock

        resp = client._get_operation_statuses()
        self.assertEqual(len(list(resp.operation_statuses)), 0)

    def test_get_operations_statuses_with_operation_extensions(self):
        """Test get_operations_statuses passes operation-level extensions and echoes them back per operation."""
        from google.protobuf import any_pb2, wrappers_pb2

        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock

        op_ext = any_pb2.Any()
        op_ext.Pack(wrappers_pb2.StringValue(value="op_extension"))

        resp = client._get_operation_statuses(
            operation_ids=["default-op-1", "default-op-2"],
            operation_extensions=[op_ext],
        )
        result = list(resp.operation_statuses)
        self.assertEqual(len(result), 2)
        self.assertEqual({s.operation_id for s in result}, {"default-op-1", "default-op-2"})

        # Verify operation-level extensions were included in the request
        self.assertEqual(len(mock.req.operation_status.extensions), 1)
        unpacked = wrappers_pb2.StringValue()
        mock.req.operation_status.extensions[0].Unpack(unpacked)
        self.assertEqual(unpacked.value, "op_extension")

        # Verify operation-level extensions were echoed back per operation
        for op_status in result:
            self.assertEqual(len(op_status.extensions), 1)
            echoed = wrappers_pb2.StringValue()
            op_status.extensions[0].Unpack(echoed)
            self.assertEqual(echoed.value, "op_extension")

    def test_get_operations_statuses_with_request_extensions(self):
        """Test _get_operation_statuses sends request-level extensions and echoes them back."""
        from google.protobuf import any_pb2, wrappers_pb2

        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
        mock = MockService(client._session_id)
        client._stub = mock

        req_ext = any_pb2.Any()
        req_ext.Pack(wrappers_pb2.StringValue(value="request_extension"))

        resp = client._get_operation_statuses(
            operation_ids=["default-op-1"],
            request_extensions=[req_ext],
        )

        # Verify the operation status is returned
        result = list(resp.operation_statuses)
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].operation_id, "default-op-1")

        # Verify request-level extensions were included in the request
        self.assertEqual(len(mock.req.extensions), 1)
        unpacked = wrappers_pb2.StringValue()
        mock.req.extensions[0].Unpack(unpacked)
        self.assertEqual(unpacked.value, "request_extension")

        # Verify request-level extensions were echoed back in the response
        self.assertEqual(len(resp.extensions), 1)
        resp_echoed = wrappers_pb2.StringValue()
        resp.extensions[0].Unpack(resp_echoed)
        self.assertEqual(resp_echoed.value, "request_extension")

    def test_analyze_plan_short_deadline_fires_then_succeeds_after_disabling(self):
        """With a short deadline the call fails; after disabling deadlines it succeeds."""

        class CapturingMock(MockService):
            """Captures the timeout passed by the client; raises DEADLINE_EXCEEDED if set."""

            def __init__(self, session_id):
                super().__init__(session_id)
                self.captured_timeout = "not_called"

            def AnalyzePlan(self, req, metadata, timeout=None):
                self.captured_timeout = timeout
                if timeout is not None:
                    raise TestException("deadline exceeded", grpc.StatusCode.DEADLINE_EXCEEDED)
                return super().AnalyzePlan(req, metadata, timeout=timeout)

        client_with_deadline = SparkConnectClient(
            "sc://foo/",
            use_reattachable_execute=False,
            rpc_deadlines=RpcDeadlines(analyze_plan=0.050),
            retry_policy=dict(max_retries=0),
        )
        mock_with_deadline = CapturingMock(session_id=client_with_deadline._session_id)
        client_with_deadline._stub = mock_with_deadline
        with self.assertRaises(SparkConnectGrpcException) as cm:
            client_with_deadline._analyze("schema", plan=proto.Plan())
        self.assertEqual(cm.exception.getGrpcStatusCode(), grpc.StatusCode.DEADLINE_EXCEEDED)
        self.assertEqual(mock_with_deadline.captured_timeout, 0.050)

        client_disabled = SparkConnectClient(
            "sc://foo/",
            use_reattachable_execute=False,
            rpc_deadlines=RpcDeadlines.disabled(),
            retry_policy=dict(max_retries=0),
        )
        mock_disabled = CapturingMock(session_id=client_disabled._session_id)
        client_disabled._stub = mock_disabled
        client_disabled._analyze("schema", plan=proto.Plan())
        self.assertIsNone(mock_disabled.captured_timeout)

    def test_each_rpc_receives_configured_deadline(self):
        """Every RPC that accepts a deadline should forward it as timeout to the stub."""

        class TimeoutCapturingMock(MockService):
            """Records the timeout kwarg for each RPC call."""

            def __init__(self, session_id):
                super().__init__(session_id)
                self.captured_timeouts = {}

            def AnalyzePlan(self, req, metadata, timeout=None):
                self.captured_timeouts["AnalyzePlan"] = timeout
                return super().AnalyzePlan(req, metadata, timeout=timeout)

            def Config(self, req, metadata, timeout=None):
                self.captured_timeouts["Config"] = timeout
                return super().Config(req, metadata, timeout=timeout)

            def Interrupt(self, req, metadata, timeout=None):
                self.captured_timeouts["Interrupt"] = timeout
                return super().Interrupt(req, metadata, timeout=timeout)

            def ReleaseSession(self, req, metadata, timeout=None):
                self.captured_timeouts["ReleaseSession"] = timeout
                resp = proto.ReleaseSessionResponse()
                resp.session_id = self._session_id
                return resp

            def GetStatus(self, req, metadata, timeout=None):
                self.captured_timeouts["GetStatus"] = timeout
                return super().GetStatus(req, metadata, timeout=timeout)

        deadlines = RpcDeadlines(
            analyze_plan=11.0,
            config=22.0,
            interrupt=33.0,
            release_session=44.0,
            get_status=55.0,
        )
        client = SparkConnectClient(
            "sc://foo/",
            use_reattachable_execute=False,
            rpc_deadlines=deadlines,
            retry_policy=dict(max_retries=0),
        )
        mock = TimeoutCapturingMock(session_id=client._session_id)
        client._stub = mock

        client._analyze("schema", plan=proto.Plan())
        self.assertEqual(mock.captured_timeouts["AnalyzePlan"], 11.0)

        op = proto.ConfigRequest.Operation()
        op.get.keys.append("spark.sql.shuffle.partitions")
        client.config(op)
        self.assertEqual(mock.captured_timeouts["Config"], 22.0)

        client.interrupt_all()
        self.assertEqual(mock.captured_timeouts["Interrupt"], 33.0)

        client.release_session()
        self.assertEqual(mock.captured_timeouts["ReleaseSession"], 44.0)

        client._get_operation_statuses()
        self.assertEqual(mock.captured_timeouts["GetStatus"], 55.0)


@unittest.skipIf(not should_test_connect, connect_requirement_message)
class SparkConnectClientReattachTestCase(unittest.TestCase):
    def setUp(self) -> None:
        self.request = proto.ExecutePlanRequest()
        self.retrying = lambda: Retrying(TestPolicy())
        self.response = proto.ExecutePlanResponse(
            response_id="1",
        )
        self.finished = proto.ExecutePlanResponse(
            result_complete=proto.ExecutePlanResponse.ResultComplete(),
            response_id="2",
        )

    def _stub_with(self, execute=None, attach=None):
        return MockSparkConnectStub(
            execute_ops=ResponseGenerator(execute) if execute is not None else None,
            attach_ops=ResponseGenerator(attach) if attach is not None else None,
        )

    def test_basic_flow(self):
        stub = self._stub_with([self.response, self.finished])
        ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])
        for b in ite:
            pass

        def check_all():
            self.assertEqual(0, stub.attach_calls)
            self.assertEqual(1, stub.release_until_calls)
            self.assertEqual(1, stub.release_calls)
            self.assertEqual(1, stub.execute_calls)

        eventually(timeout=1, catch_assertions=True)(check_all)()

    def test_fail_during_execute(self):
        def fatal():
            raise TestException("Fatal")

        stub = self._stub_with([self.response, fatal])
        with self.assertRaises(TestException):
            ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])
            for b in ite:
                pass

        def check():
            self.assertEqual(0, stub.attach_calls)
            self.assertEqual(1, stub.release_calls)
            self.assertEqual(1, stub.release_until_calls)
            self.assertEqual(1, stub.execute_calls)

        eventually(timeout=1, catch_assertions=True)(check)()

    def test_fail_and_retry_during_execute(self):
        def non_fatal():
            raise TestException("Non Fatal", grpc.StatusCode.UNAVAILABLE)

        stub = self._stub_with(
            [self.response, non_fatal], [self.response, self.response, self.finished]
        )
        ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])
        for b in ite:
            pass

        def check():
            self.assertEqual(1, stub.attach_calls)
            self.assertEqual(1, stub.release_calls)
            self.assertEqual(3, stub.release_until_calls)
            self.assertEqual(1, stub.execute_calls)

        eventually(timeout=1, catch_assertions=True)(check)()

    def test_fail_and_retry_during_reattach(self):
        count = 0

        def non_fatal():
            nonlocal count
            if count < 2:
                count += 1
                raise TestException("Non Fatal", grpc.StatusCode.UNAVAILABLE)
            else:
                return proto.ExecutePlanResponse()

        stub = self._stub_with(
            [self.response, non_fatal], [self.response, non_fatal, self.response, self.finished]
        )
        ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])
        for b in ite:
            pass

        def check():
            self.assertEqual(2, stub.attach_calls)
            self.assertEqual(3, stub.release_until_calls)
            self.assertEqual(1, stub.release_calls)
            self.assertEqual(1, stub.execute_calls)

        eventually(timeout=1, catch_assertions=True)(check)()

    def test_not_found_recovers(self):
        """SPARK-48056: Assert that the client recovers from session or operation not
        found error if no partial responses were previously received.
        """

        def not_found_recovers(error_code: str):
            def not_found():
                raise TestException(
                    error_code,
                    grpc.StatusCode.UNAVAILABLE,
                    trailing_status=status_pb2.Status(code=14, message=error_code, details=""),
                )

            stub = self._stub_with([not_found, self.finished])
            ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])

            for _ in ite:
                pass

            def checks():
                self.assertEqual(2, stub.execute_calls)
                self.assertEqual(0, stub.attach_calls)
                self.assertEqual(1, stub.release_calls)
                self.assertEqual(0, stub.release_until_calls)

            eventually(timeout=1, catch_assertions=True)(checks)()

        parameters = ["INVALID_HANDLE.SESSION_NOT_FOUND", "INVALID_HANDLE.OPERATION_NOT_FOUND"]
        for b in parameters:
            not_found_recovers(b)

    def test_not_found_fails(self):
        """SPARK-48056: Assert that the client fails from session or operation not found error
        if a partial response was previously received.
        """

        def not_found_fails(error_code: str):
            def not_found():
                raise TestException(
                    error_code,
                    grpc.StatusCode.UNAVAILABLE,
                    trailing_status=status_pb2.Status(code=14, message=error_code, details=""),
                )

            stub = self._stub_with([self.response], [not_found])

            with self.assertRaises(PySparkRuntimeError) as e:
                ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])
                for _ in ite:
                    pass

            self.assertTrue("RESPONSE_ALREADY_RECEIVED" in e.exception.getMessage())
            self.assertTrue(error_code in e.exception.getMessage())

            def checks():
                self.assertEqual(1, stub.execute_calls)
                self.assertEqual(1, stub.attach_calls)
                self.assertEqual(1, stub.release_calls)
                self.assertEqual(1, stub.release_until_calls)

            eventually(timeout=1, catch_assertions=True)(checks)()

        parameters = ["INVALID_HANDLE.SESSION_NOT_FOUND", "INVALID_HANDLE.OPERATION_NOT_FOUND"]
        for b in parameters:
            not_found_fails(b)

    def test_observed_session_id(self):
        stub = self._stub_with([self.response, self.finished])
        ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])
        session_id = "test-session-id"

        reattach = ite._create_reattach_execute_request()
        self.assertEqual(reattach.client_observed_server_side_session_id, "")

        self.request.client_observed_server_side_session_id = session_id
        reattach = ite._create_reattach_execute_request()
        self.assertEqual(reattach.client_observed_server_side_session_id, session_id)

    def test_deadline_exceeded_triggers_reattach(self):
        """DEADLINE_EXCEEDED mid-stream on ExecutePlan should trigger a ReattachExecute."""

        def deadline_exceeded():
            raise TestException("deadline", grpc.StatusCode.DEADLINE_EXCEEDED)

        stub = self._stub_with(
            [self.response, deadline_exceeded],
            [self.response, self.finished],
        )
        ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])
        for _ in ite:
            pass

        def check():
            self.assertEqual(1, stub.execute_calls)
            self.assertEqual(1, stub.attach_calls)
            self.assertEqual(1, stub.release_calls)

        eventually(timeout=1, catch_assertions=True)(check)()

    def test_deadline_exceeded_mid_stream_completes_successfully(self):
        """After a mid-stream DEADLINE_EXCEEDED, reattach resumes and all responses are collected."""

        response2 = proto.ExecutePlanResponse(response_id="2")
        response3 = proto.ExecutePlanResponse(response_id="3")

        def deadline_exceeded():
            raise TestException("deadline", grpc.StatusCode.DEADLINE_EXCEEDED)

        finished = proto.ExecutePlanResponse(
            result_complete=proto.ExecutePlanResponse.ResultComplete(),
            response_id="final",
        )
        stub = self._stub_with(
            [self.response, response2, deadline_exceeded],
            [response3, finished],
        )
        collected = []
        ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, [])
        for r in ite:
            if not r.HasField("result_complete"):
                collected.append(r.response_id)

        self.assertEqual(collected, ["1", "2", "3"])

        def check():
            self.assertEqual(1, stub.execute_calls)
            self.assertEqual(1, stub.attach_calls)
            self.assertEqual(1, stub.release_calls)

        eventually(timeout=1, catch_assertions=True)(check)()

    def test_server_unreachable(self):
        # DNS resolution should fail for "foo". This error is a retriable UNAVAILABLE error.
        client = SparkConnectClient(
            "sc://foo", use_reattachable_execute=False, retry_policy=dict(max_retries=0)
        )
        with self.assertRaises(SparkConnectGrpcException) as cm:
            command = proto.Command()
            client.execute_command(command)
        err = cm.exception
        self.assertEqual(err.getGrpcStatusCode(), grpc.StatusCode.UNAVAILABLE)
        self.assertEqual(err.getErrorClass(), None)
        self.assertEqual(err.getSqlState(), None)

    def test_error_codes(self):
        msg = "Something went wrong on the server"

        def raise_without_status():
            raise TestException(msg=msg, trailing_status=None)

        def raise_without_status_unauthenticated():
            raise TestException(msg=msg, code=grpc.StatusCode.UNAUTHENTICATED)

        def raise_without_status_permission_denied():
            raise TestException(msg=msg, code=grpc.StatusCode.PERMISSION_DENIED)

        def raise_without_details():
            status = status_pb2.Status(
                code=grpc.StatusCode.INTERNAL.value[0], message=msg, details=[]
            )
            raise TestException(msg=msg, trailing_status=status)

        def raise_without_metadata():
            any = any_pb2.Any()
            any.Pack(ErrorInfo())
            status = status_pb2.Status(
                code=grpc.StatusCode.INTERNAL.value[0], message=msg, details=[any]
            )
            raise TestException(msg=msg, trailing_status=status)

        def raise_with_error_class():
            any = any_pb2.Any()
            any.Pack(ErrorInfo(metadata=dict(errorClass="TEST_ERROR_CLASS")))
            status = status_pb2.Status(
                code=grpc.StatusCode.INTERNAL.value[0], message=msg, details=[any]
            )
            raise TestException(msg=msg, trailing_status=status)

        def raise_with_sql_state():
            any = any_pb2.Any()
            any.Pack(ErrorInfo(metadata=dict(sqlState="TEST_SQL_STATE")))
            status = status_pb2.Status(
                code=grpc.StatusCode.INTERNAL.value[0], message=msg, details=[any]
            )
            raise TestException(msg=msg, trailing_status=status)

        test_cases = [
            (raise_without_status, grpc.StatusCode.INTERNAL, None, None),
            (raise_without_status_unauthenticated, grpc.StatusCode.UNAUTHENTICATED, None, None),
            (raise_without_status_permission_denied, grpc.StatusCode.PERMISSION_DENIED, None, None),
            (raise_without_details, grpc.StatusCode.INTERNAL, None, None),
            (raise_without_metadata, grpc.StatusCode.INTERNAL, None, None),
            (raise_with_error_class, grpc.StatusCode.INTERNAL, "TEST_ERROR_CLASS", None),
            (raise_with_sql_state, grpc.StatusCode.INTERNAL, None, "TEST_SQL_STATE"),
        ]

        for (
            response_function,
            expected_status_code,
            expected_error_class,
            expected_sql_state,
        ) in test_cases:
            client = SparkConnectClient(
                "sc://foo", use_reattachable_execute=False, retry_policy=dict(max_retries=0)
            )
            client._stub = self._stub_with([response_function])
            with self.assertRaises(SparkConnectGrpcException) as cm:
                command = proto.Command()
                client.execute_command(command)
            err = cm.exception
            self.assertEqual(err.getGrpcStatusCode(), expected_status_code)
            self.assertEqual(err.getErrorClass(), expected_error_class)
            self.assertEqual(err.getSqlState(), expected_sql_state)


if __name__ == "__main__":
    from pyspark.testing import main

    main()