core.py

#
# 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.
#
__all__ = [
    "ChannelBuilder",
    "DefaultChannelBuilder",
    "SparkConnectClient",
]

import atexit

import pyspark
from pyspark.sql.connect.proto.base_pb2 import FetchErrorDetailsResponse

import concurrent.futures
import logging
import threading
import os
import copy
import platform
import urllib.parse
import uuid
import sys
import time
import traceback
import weakref
from typing import (
    Iterable,
    Iterator,
    Optional,
    Any,
    Union,
    List,
    Tuple,
    Dict,
    Set,
    NoReturn,
    Mapping,
    cast,
    TYPE_CHECKING,
    Type,
)

import pandas as pd
import pyarrow as pa

import google.protobuf.message
from grpc_status import rpc_status
import grpc
from google.protobuf import text_format, any_pb2
from google.rpc import error_details_pb2

from pyspark.util import is_remote_only, disable_gc
from pyspark.accumulators import SpecialAccumulatorIds, pickleSer
from pyspark.version import __version__
from pyspark.traceback_utils import CallSite
from pyspark.resource.information import ResourceInformation
from pyspark.sql.metrics import MetricValue, PlanMetrics, ExecutionInfo, ObservedMetrics
from pyspark.sql.connect.client.artifact import ArtifactManager
from pyspark.sql.connect.logging import logger
from pyspark.sql.connect.profiler import ConnectProfilerCollector
from pyspark.sql.connect.client.reattach import ExecutePlanResponseReattachableIterator
from pyspark.sql.connect.client.retries import RetryPolicy, Retrying, DefaultPolicy
from pyspark.sql.connect.conversion import (
    storage_level_to_proto,
    proto_to_storage_level,
    proto_to_remote_cached_dataframe,
)
import pyspark.sql.connect.proto as pb2
import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib
import pyspark.sql.connect.types as types
from pyspark.errors.exceptions.connect import (
    convert_exception,
    convert_observation_errors,
    SparkConnectException,
    SparkConnectGrpcException,
)
from pyspark.sql.connect.expressions import (
    LiteralExpression,
    PythonUDF,
    CommonInlineUserDefinedFunction,
    JavaUDF,
)
from pyspark.sql.connect.plan import (
    CommonInlineUserDefinedTableFunction,
    CommonInlineUserDefinedDataSource,
    PythonUDTF,
    PythonDataSource,
)
from pyspark.sql.connect.observation import Observation
from pyspark.sql.connect.utils import get_python_ver
from pyspark.sql.pandas.types import from_arrow_schema
from pyspark.sql.pandas.conversion import _convert_arrow_table_to_pandas
from pyspark.sql.types import DataType, StructType
from pyspark.util import PythonEvalType
from pyspark.storagelevel import StorageLevel
from pyspark.errors import PySparkValueError, PySparkAssertionError, PySparkNotImplementedError
from pyspark.sql.connect.shell.progress import Progress, ProgressHandler, from_proto

if TYPE_CHECKING:
    from google.rpc.error_details_pb2 import ErrorInfo
    from google.rpc.status_pb2 import Status
    from pyspark.sql.connect._typing import DataTypeOrString
    from pyspark.sql.connect.session import SparkSession
    from pyspark.sql.datasource import DataSource


PYSPARK_ROOT = os.path.dirname(pyspark.__file__)


def _import_zstandard_if_available() -> Optional[Any]:
    """
    Import zstandard if available, otherwise return None.
    This is used to handle the case when zstandard is not installed.
    """
    try:
        import zstandard

        return zstandard
    except ImportError:
        return None


class ChannelBuilder:
    """
    This is a helper class that is used to create a GRPC channel based on the given
    connection string per the documentation of Spark Connect.

    The standard implementation is in :class:`DefaultChannelBuilder`.
    """

    PARAM_USE_SSL = "use_ssl"
    PARAM_TOKEN = "token"
    PARAM_USER_ID = "user_id"
    PARAM_USER_AGENT = "user_agent"
    PARAM_SESSION_ID = "session_id"

    GRPC_MAX_MESSAGE_LENGTH_DEFAULT = 128 * 1024 * 1024

    GRPC_DEFAULT_OPTIONS = [
        ("grpc.max_send_message_length", GRPC_MAX_MESSAGE_LENGTH_DEFAULT),
        ("grpc.max_receive_message_length", GRPC_MAX_MESSAGE_LENGTH_DEFAULT),
    ]

    def __init__(
        self,
        channelOptions: Optional[List[Tuple[str, Any]]] = None,
        params: Optional[Dict[str, str]] = None,
    ):
        self._interceptors: List[grpc.UnaryStreamClientInterceptor] = []
        self._params: Dict[str, str] = params or dict()
        self._channel_options: List[Tuple[str, Any]] = ChannelBuilder.GRPC_DEFAULT_OPTIONS.copy()

        if channelOptions is not None:
            for key, value in channelOptions:
                self.setChannelOption(key, value)

    def get(self, key: str) -> Any:
        """
        Parameters
        ----------
        key : str
            Parameter key name.
        Returns
        -------
        The parameter value if present, raises exception otherwise.
        """
        return self._params[key]

    def getDefault(self, key: str, default: Any) -> Any:
        return self._params.get(key, default)

    def set(self, key: str, value: Any) -> None:
        self._params[key] = value

    def setChannelOption(self, key: str, value: Any) -> None:
        # overwrite option if it exists already else append it
        for i, option in enumerate(self._channel_options):
            if option[0] == key:
                self._channel_options[i] = (key, value)
                return
        self._channel_options.append((key, value))

    def add_interceptor(self, interceptor: grpc.UnaryStreamClientInterceptor) -> None:
        self._interceptors.append(interceptor)

    def toChannel(self) -> grpc.Channel:
        """
        The actual channel builder implementations should implement this function
        to return grpc Channel.
        This function should generally use self._insecure_channel or
        self._secure_channel so that configuration options are applied
        appropriately.
        """
        raise PySparkNotImplementedError

    @property
    def host(self) -> str:
        """
        The hostname where this client intends to connect.
        This is used for end-user display purpose in REPL
        """
        raise PySparkNotImplementedError

    def _insecure_channel(self, target: Any, **kwargs: Any) -> grpc.Channel:
        channel = grpc.insecure_channel(target, options=self._channel_options, **kwargs)

        if len(self._interceptors) > 0:
            logger.debug(f"Applying interceptors ({self._interceptors})")
            channel = grpc.intercept_channel(channel, *self._interceptors)
        return channel

    def _secure_channel(self, target: Any, credentials: Any, **kwargs: Any) -> grpc.Channel:
        channel = grpc.secure_channel(target, credentials, options=self._channel_options, **kwargs)

        if len(self._interceptors) > 0:
            logger.debug(f"Applying interceptors ({self._interceptors})")
            channel = grpc.intercept_channel(channel, *self._interceptors)
        return channel

    @property
    def userId(self) -> Optional[str]:
        """
        Returns
        -------
        The user_id (extracted from connection string or configured by other means).
        """
        return self._params.get(ChannelBuilder.PARAM_USER_ID, None)

    @property
    def token(self) -> Optional[str]:
        return self._params.get(
            ChannelBuilder.PARAM_TOKEN, os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN")
        )

    def metadata(self) -> Iterable[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
        by the channel.
        Returns
        -------
        A list of tuples (key, value)
        """
        return [
            (k, self._params[k])
            for k in self._params
            if k
            not in [
                ChannelBuilder.PARAM_TOKEN,
                ChannelBuilder.PARAM_USE_SSL,
                ChannelBuilder.PARAM_USER_ID,
                ChannelBuilder.PARAM_USER_AGENT,
                ChannelBuilder.PARAM_SESSION_ID,
            ]
        ]

    @property
    def session_id(self) -> Optional[str]:
        """
        Returns
        -------
        The session_id extracted from the parameters of the connection string or `None` if not
        specified.
        """
        session_id = self._params.get(ChannelBuilder.PARAM_SESSION_ID, None)
        if session_id is not None:
            try:
                uuid.UUID(session_id, version=4)
            except ValueError as ve:
                raise PySparkValueError(
                    errorClass="INVALID_SESSION_UUID_ID",
                    messageParameters={"arg_name": "session_id", "origin": str(ve)},
                )
        return session_id

    @property
    def userAgent(self) -> str:
        """
        Returns
        -------
        user_agent : str
            The user_agent parameter specified in the connection string,
            or "_SPARK_CONNECT_PYTHON" when not specified.
            The returned value will be percent encoded.
        """
        user_agent = self._params.get(
            ChannelBuilder.PARAM_USER_AGENT,
            os.getenv("SPARK_CONNECT_USER_AGENT", "_SPARK_CONNECT_PYTHON"),
        )

        ua_len = len(urllib.parse.quote(user_agent))
        if ua_len > 2048:
            raise SparkConnectException(
                f"'user_agent' parameter should not exceed 2048 characters after URL "
                f"escaping, found {ua_len} characters."
            )
        return " ".join(
            [
                user_agent,
                f"spark/{__version__}",
                f"os/{platform.uname().system.lower()}",
                f"python/{platform.python_version()}",
            ]
        )


class DefaultChannelBuilder(ChannelBuilder):
    """
    This is a helper class that is used to create a GRPC channel based on the given
    connection string per the documentation of Spark Connect.

    .. versionadded:: 3.4.0

    Examples
    --------
    >>> cb =  DefaultChannelBuilder("sc://localhost")
    ... cb.endpoint
    "localhost:15002"

    >>> cb = DefaultChannelBuilder("sc://localhost/;use_ssl=true;token=aaa")
    ... cb.secure
    True
    """

    @staticmethod
    def default_port() -> int:
        if "SPARK_TESTING" in os.environ and not is_remote_only():
            from pyspark.sql.session import SparkSession as PySparkSession

            # In the case when Spark Connect uses the local mode, it starts the regular Spark
            # session that starts Spark Connect server that sets `SparkSession._instantiatedSession`
            # via SparkSession.__init__.
            #
            # We are getting the actual server port from the Spark session via Py4J to address
            # the case when the server port is set to 0 (in which allocates an ephemeral port).
            #
            # This is only used in the test/development mode.
            session = PySparkSession._instantiatedSession

            if session is not None:
                jvm = PySparkSession._instantiatedSession._jvm  # type: ignore[union-attr]
                return getattr(
                    getattr(
                        jvm,
                        "org.apache.spark.sql.connect.service.SparkConnectService$",
                    ),
                    "MODULE$",
                ).localPort()
        return 15002

    def __init__(self, url: str, channelOptions: Optional[List[Tuple[str, Any]]] = None) -> None:
        """
        Constructs a new channel builder. This is used to create the proper GRPC channel from
        the connection string.

        Parameters
        ----------
        url : str
            Spark Connect connection string
        channelOptions: list of tuple, optional
            Additional options that can be passed to the GRPC channel construction.
        """

        super().__init__(channelOptions=channelOptions)

        # Explicitly check the scheme of the URL.
        if url[:5] != "sc://":
            raise PySparkValueError(
                errorClass="INVALID_CONNECT_URL",
                messageParameters={
                    "detail": "The URL must start with 'sc://'. Please update the URL to "
                    "follow the correct format, e.g., 'sc://hostname:port'.",
                },
            )
        # Rewrite the URL to use http as the scheme so that we can leverage
        # Python's built-in parser.
        tmp_url = "http" + url[2:]
        self.url = urllib.parse.urlparse(tmp_url)
        if len(self.url.path) > 0 and self.url.path != "/":
            raise PySparkValueError(
                errorClass="INVALID_CONNECT_URL",
                messageParameters={
                    "detail": f"The path component '{self.url.path}' must be empty. Please update "
                    f"the URL to follow the correct format, e.g., 'sc://hostname:port'.",
                },
            )
        self._extract_attributes()

    def _extract_attributes(self) -> None:
        if len(self.url.params) > 0:
            parts = self.url.params.split(";")
            for p in parts:
                kv = p.split("=")
                if len(kv) != 2:
                    raise PySparkValueError(
                        errorClass="INVALID_CONNECT_URL",
                        messageParameters={
                            "detail": f"Parameter '{p}' should be provided as a "
                            f"key-value pair separated by an equal sign (=). Please update "
                            f"the parameter to follow the correct format, e.g., 'key=value'.",
                        },
                    )
                self.set(kv[0], urllib.parse.unquote(kv[1]))

        if not self.url.hostname:
            raise PySparkValueError(
                errorClass="INVALID_CONNECT_URL",
                messageParameters={
                    "detail": f"Hostname is missing in the URL: '{self.url.geturl()}'. "
                    "Please update the URL to follow the correct format, "
                    "e.g., 'sc://hostname:port'.",
                },
            )
        self._host = f"[{self.url.hostname}]" if ":" in self.url.hostname else self.url.hostname
        self._port = (
            self.url.port if self.url.port is not None else DefaultChannelBuilder.default_port()
        )

    @property
    def secure(self) -> bool:
        return self.use_ssl or self.token is not None

    @property
    def use_ssl(self) -> bool:
        return self.getDefault(ChannelBuilder.PARAM_USE_SSL, "").lower() == "true"

    @property
    def host(self) -> str:
        """
        The hostname where this client intends to connect.
        """
        return self._host

    @property
    def endpoint(self) -> str:
        return f"{self._host}:{self._port}"

    def toChannel(self) -> grpc.Channel:
        """
        Applies the parameters of the connection string and creates a new
        GRPC channel according to the configuration. Passes optional channel options to
        construct the channel.

        Returns
        -------
        GRPC Channel instance.
        """

        if not self.secure:
            return self._insecure_channel(self.endpoint)
        elif not self.use_ssl and self._host == "localhost":
            creds = grpc.local_channel_credentials()

            if self.token is not None:
                creds = grpc.composite_channel_credentials(
                    creds, grpc.access_token_call_credentials(self.token)
                )
            return self._secure_channel(self.endpoint, creds)
        else:
            creds = grpc.ssl_channel_credentials()

            if self.token is not None:
                creds = grpc.composite_channel_credentials(
                    creds, grpc.access_token_call_credentials(self.token)
                )

            return self._secure_channel(self.endpoint, creds)


class PlanObservedMetrics(ObservedMetrics):
    def __init__(self, name: str, metrics: List[pb2.Expression.Literal], keys: List[str]):
        self._name = name
        self._metrics = metrics
        self._keys = keys if keys else [f"observed_metric_{i}" for i in range(len(self.metrics))]

    def __repr__(self) -> str:
        return f"Plan observed({self._name}={self._metrics})"

    @property
    def name(self) -> str:
        return self._name

    @property
    def metrics(self) -> List[pb2.Expression.Literal]:
        return self._metrics

    @property
    def pairs(self) -> dict[str, Any]:
        result = {}
        for x in range(len(self._metrics)):
            result[self.keys[x]] = LiteralExpression._to_value(self.metrics[x])
        return result

    @property
    def keys(self) -> List[str]:
        return self._keys

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-serializable dictionary representation of this observed metrics.

        Returns
        -------
        dict
            A dictionary with keys 'name', 'keys', and 'pairs'.
        """
        return {
            "name": self._name,
            "keys": self._keys,
            "pairs": self.pairs,
        }


class AnalyzeResult:
    def __init__(
        self,
        schema: Optional[DataType],
        explain_string: Optional[str],
        tree_string: Optional[str],
        is_local: Optional[bool],
        is_streaming: Optional[bool],
        input_files: Optional[List[str]],
        spark_version: Optional[str],
        parsed: Optional[DataType],
        is_same_semantics: Optional[bool],
        semantic_hash: Optional[int],
        storage_level: Optional[StorageLevel],
        ddl_string: Optional[str],
    ):
        self.schema = schema
        self.explain_string = explain_string
        self.tree_string = tree_string
        self.is_local = is_local
        self.is_streaming = is_streaming
        self.input_files = input_files
        self.spark_version = spark_version
        self.parsed = parsed
        self.is_same_semantics = is_same_semantics
        self.semantic_hash = semantic_hash
        self.storage_level = storage_level
        self.ddl_string = ddl_string

    @classmethod
    def fromProto(cls, pb: Any) -> "AnalyzeResult":
        schema: Optional[DataType] = None
        explain_string: Optional[str] = None
        tree_string: Optional[str] = None
        is_local: Optional[bool] = None
        is_streaming: Optional[bool] = None
        input_files: Optional[List[str]] = None
        spark_version: Optional[str] = None
        parsed: Optional[DataType] = None
        is_same_semantics: Optional[bool] = None
        semantic_hash: Optional[int] = None
        storage_level: Optional[StorageLevel] = None
        ddl_string: Optional[str] = None

        if pb.HasField("schema"):
            schema = types.proto_schema_to_pyspark_data_type(pb.schema.schema)
        elif pb.HasField("explain"):
            explain_string = pb.explain.explain_string
        elif pb.HasField("tree_string"):
            tree_string = pb.tree_string.tree_string
        elif pb.HasField("is_local"):
            is_local = pb.is_local.is_local
        elif pb.HasField("is_streaming"):
            is_streaming = pb.is_streaming.is_streaming
        elif pb.HasField("input_files"):
            input_files = pb.input_files.files
        elif pb.HasField("spark_version"):
            spark_version = pb.spark_version.version
        elif pb.HasField("ddl_parse"):
            parsed = types.proto_schema_to_pyspark_data_type(pb.ddl_parse.parsed)
        elif pb.HasField("same_semantics"):
            is_same_semantics = pb.same_semantics.result
        elif pb.HasField("semantic_hash"):
            semantic_hash = pb.semantic_hash.result
        elif pb.HasField("persist"):
            pass
        elif pb.HasField("unpersist"):
            pass
        elif pb.HasField("get_storage_level"):
            storage_level = proto_to_storage_level(pb.get_storage_level.storage_level)
        elif pb.HasField("json_to_ddl"):
            ddl_string = pb.json_to_ddl.ddl_string
        else:
            raise SparkConnectException("No analyze result found!")

        return AnalyzeResult(
            schema,
            explain_string,
            tree_string,
            is_local,
            is_streaming,
            input_files,
            spark_version,
            parsed,
            is_same_semantics,
            semantic_hash,
            storage_level,
            ddl_string,
        )


class ConfigResult:
    def __init__(self, pairs: List[Tuple[str, Optional[str]]], warnings: List[str]):
        self.pairs = pairs
        self.warnings = warnings

    @classmethod
    def fromProto(cls, pb: pb2.ConfigResponse) -> "ConfigResult":
        return ConfigResult(
            pairs=[(pair.key, pair.value if pair.HasField("value") else None) for pair in pb.pairs],
            warnings=list(pb.warnings),
        )


def _is_pyspark_source(filename: str) -> bool:
    """Check if the given filename is from the pyspark package."""
    return filename.startswith(PYSPARK_ROOT)


class SparkConnectClient(object):
    """
    Conceptually the remote spark session that communicates with the server
    """

    def __init__(
        self,
        connection: Union[str, ChannelBuilder],
        user_id: Optional[str] = None,
        channel_options: Optional[List[Tuple[str, Any]]] = None,
        retry_policy: Optional[Dict[str, Any]] = None,
        use_reattachable_execute: bool = True,
        session_hooks: Optional[list["SparkSession.Hook"]] = None,
        allow_arrow_batch_chunking: bool = True,
        preferred_arrow_chunk_size: Optional[int] = None,
    ):
        """
        Creates a new SparkSession for the Spark Connect interface.

        Parameters
        ----------
        connection : str or :class:`ChannelBuilder`
            Connection string that is used to extract the connection parameters and configure
            the GRPC connection. Or instance of ChannelBuilder that creates GRPC connection.
            Defaults to `sc://localhost`.
        user_id : str, optional
            Optional unique user ID that is used to differentiate multiple users and
            isolate their Spark Sessions. If the `user_id` is not set, will default to
            the $USER environment. Defining the user ID as part of the connection string
            takes precedence.
        channel_options: list of tuple, optional
            Additional options that can be passed to the GRPC channel construction.
        retry_policy: dict of str and any, optional
            Additional configuration for retrying. There are four configurations as below
                * ``max_retries``
                    Maximum number of tries default 15
                * ``backoff_multiplier``
                    Backoff multiplier for the policy. Default: 4(ms)
                * ``initial_backoff``
                    Backoff to wait before the first retry. Default: 50(ms)
                * ``max_backoff``
                    Maximum backoff controls the maximum amount of time to wait before retrying
                    a failed request. Default: 60000(ms).
        use_reattachable_execute: bool
            Enable reattachable execution.
        session_hooks: list[SparkSession.Hook], optional
            List of session hooks to call.
        allow_arrow_batch_chunking: bool
            Whether to allow the server to split large Arrow batches into smaller chunks.
            Although Arrow results are split into batches with a size limit according to estimation,
            the size of the batches is not guaranteed to be less than the limit, especially when a
            single row is larger than the limit, in which case the server will fail to split it
            further into smaller batches. As a result, the client may encounter a gRPC error stating
            "Received message larger than max" when a batch is too large.
            If true, the server will split large Arrow batches into smaller chunks, and the client
            is expected to handle the chunked Arrow batches.
            If false, the server will not chunk large Arrow batches.
        preferred_arrow_chunk_size: Optional[int]
            Optional preferred Arrow batch size in bytes for the server to use when sending Arrow
            results.
            The server will attempt to use this size if it is set and within the valid range
            ([1KB, max batch size on server]). Otherwise, the server's maximum batch size is used.
        """
        self.thread_local = threading.local()

        # Parse the connection string.
        self._builder = (
            connection
            if isinstance(connection, ChannelBuilder)
            else DefaultChannelBuilder(connection, channel_options)
        )
        self._user_id = None
        self._retry_policies: List[RetryPolicy] = []

        retry_policy_args = retry_policy or dict()
        default_policy = DefaultPolicy(**retry_policy_args)
        self.set_retry_policies([default_policy])

        if self._builder.session_id is None:
            # Generate a unique session ID for this client. This UUID must be unique to allow
            # concurrent Spark sessions of the same user. If the channel is closed, creating
            # a new client will create a new session ID.
            self._session_id = str(uuid.uuid4())
        else:
            # Use the pre-defined session ID.
            self._session_id = str(self._builder.session_id)

        if self._builder.userId is not None:
            self._user_id = self._builder.userId
        elif user_id is not None:
            self._user_id = user_id
        else:
            self._user_id = os.getenv("SPARK_USER", os.getenv("USER", None))

        self._channel = self._builder.toChannel()
        self._closed = False
        self._internal_stub = grpc_lib.SparkConnectServiceStub(self._channel)
        self._artifact_manager = ArtifactManager(
            self._user_id, self._session_id, self._channel, self._builder.metadata()
        )
        self._use_reattachable_execute = use_reattachable_execute
        self._allow_arrow_batch_chunking = allow_arrow_batch_chunking
        self._preferred_arrow_chunk_size = preferred_arrow_chunk_size
        self._session_hooks = session_hooks or []
        # Configure logging for the SparkConnect client.

        # Capture the server-side session ID and set it to None initially. It will
        # be updated on the first response received.
        self._server_session_id: Optional[str] = None

        self._profiler_collector = ConnectProfilerCollector()

        self._progress_handlers: List[ProgressHandler] = []
        self._execution_info_callbacks: "List[Callable[[str, ExecutionInfo], None]]" = []

        self._zstd_module = _import_zstandard_if_available()
        self._plan_compression_threshold: Optional[int] = None  # Will be fetched lazily
        self._plan_compression_algorithm: Optional[str] = None  # Will be fetched lazily

        self._release_futures: weakref.WeakSet[concurrent.futures.Future] = weakref.WeakSet()

        self._release_session_on_exit = os.getenv(
            "SPARK_CONNECT_RELEASE_SESSION_ON_EXIT", "false"
        ).lower() in ("true", "1")
        # cleanup if possible
        atexit.register(self._on_exit)

        self.global_user_context_extensions: List[Tuple[str, any_pb2.Any]] = []
        self.global_user_context_extensions_lock = threading.Lock()

    @property
    def _stub(self) -> grpc_lib.SparkConnectServiceStub:
        if self.is_closed:
            raise SparkConnectException(
                errorClass="NO_ACTIVE_SESSION", messageParameters=dict()
            ) from None
        return self._internal_stub

    # For testing only.
    @_stub.setter
    def _stub(self, value: grpc_lib.SparkConnectServiceStub) -> None:
        self._internal_stub = value

    def register_progress_handler(self, handler: ProgressHandler) -> None:
        """
        Register a progress handler to be called when a progress message is received.

        Parameters
        ----------
        handler : ProgressHandler
          The callable that will be called with the progress information.

        """
        if handler in self._progress_handlers:
            return
        self._progress_handlers.append(handler)

    def clear_progress_handlers(self) -> None:
        self._progress_handlers.clear()

    def register_execution_info_callback(
        self, cb: "Callable[[str, ExecutionInfo], None]"
    ) -> None:
        if cb not in self._execution_info_callbacks:
            self._execution_info_callbacks.append(cb)

    def remove_execution_info_callback(
        self, cb: "Callable[[str, ExecutionInfo], None]"
    ) -> None:
        if cb in self._execution_info_callbacks:
            self._execution_info_callbacks.remove(cb)

    def _fire_execution_info(self, operation_id: str, ei: ExecutionInfo) -> None:
        for cb in self._execution_info_callbacks:
            try:
                cb(operation_id, ei)
            except Exception:
                pass

    def remove_progress_handler(self, handler: ProgressHandler) -> None:
        """
        Remove a progress handler from the list of registered handlers.

        Parameters
        ----------
        handler : ProgressHandler
          The callable to remove from the list of progress handlers.

        """
        self._progress_handlers.remove(handler)

    def _retrying(self) -> "Retrying":
        return Retrying(self._retry_policies)

    def disable_reattachable_execute(self) -> "SparkConnectClient":
        self._use_reattachable_execute = False
        return self

    def enable_reattachable_execute(self) -> "SparkConnectClient":
        self._use_reattachable_execute = True
        return self

    def set_retry_policies(self, policies: Iterable[RetryPolicy]) -> None:
        """
        Sets list of policies to be used for retries.
        I.e. set_retry_policies([DefaultPolicy(), CustomPolicy()]).

        """
        self._retry_policies = list(policies)

    def get_retry_policies(self) -> List[RetryPolicy]:
        """
        Return list of currently used policies
        """
        return list(self._retry_policies)

    @classmethod
    def _retrieve_stack_frames(cls) -> List[CallSite]:
        """
        Return a list of CallSites representing the relevant stack frames in the callstack.
        """
        frames = traceback.extract_stack()

        filtered_stack_frames = []
        for i, frame in enumerate(frames):
            filename, lineno, func, _ = frame
            if _is_pyspark_source(filename):
                # Do not include PySpark internal frames as they are not user application code
                break
            if i + 1 < len(frames):
                _, _, func, _ = frames[i + 1]
            filtered_stack_frames.append(CallSite(function=func, file=filename, linenum=lineno))

        return filtered_stack_frames

    @classmethod
    def _build_call_stack_trace(cls) -> Optional[any_pb2.Any]:
        """
        Build a call stack trace for the current Spark Connect action
        Returns
        -------
        FetchErrorDetailsResponse.Error: An Error object containing list of stack frames
        of the user code packed as Any protobuf
        """
        if os.getenv("SPARK_CONNECT_DEBUG_CLIENT_CALL_STACK", "false").lower() in ("true", "1"):
            stack_frames = cls._retrieve_stack_frames()
            call_stack = FetchErrorDetailsResponse.Error()
            for call_site in stack_frames:
                stack_trace_element = pb2.FetchErrorDetailsResponse.StackTraceElement()
                stack_trace_element.declaring_class = ""  # unknown information
                stack_trace_element.method_name = call_site.function
                stack_trace_element.file_name = call_site.file
                stack_trace_element.line_number = call_site.linenum
                call_stack.stack_trace.append(stack_trace_element)
            if len(call_stack.stack_trace) > 0:
                call_stack_details = any_pb2.Any()
                call_stack_details.Pack(call_stack)
                return call_stack_details
        return None

    def register_udf(
        self,
        function: Any,
        return_type: "DataTypeOrString",
        name: Optional[str] = None,
        eval_type: int = PythonEvalType.SQL_BATCHED_UDF,
        deterministic: bool = True,
    ) -> str:
        """
        Create a temporary UDF in the session catalog on the other side. We generate a
        temporary name for it.
        """

        if name is None:
            name = f"fun_{uuid.uuid4().hex}"

        # construct a PythonUDF
        py_udf = PythonUDF(
            output_type=return_type,
            eval_type=eval_type,
            func=function,
            python_ver="%d.%d" % sys.version_info[:2],
        )

        # construct a CommonInlineUserDefinedFunction
        fun = CommonInlineUserDefinedFunction(
            function_name=name,
            arguments=[],
            function=py_udf,
            deterministic=deterministic,
        ).to_plan_udf(self)

        # construct the request
        req = self._execute_plan_request_with_metadata()
        req.plan.command.register_function.CopyFrom(fun)

        self._execute(req)
        return name

    def register_udtf(
        self,
        function: Any,
        return_type: Optional["DataTypeOrString"],
        name: str,
        eval_type: int = PythonEvalType.SQL_TABLE_UDF,
        deterministic: bool = True,
    ) -> str:
        """
        Register a user-defined table function (UDTF) in the session catalog
        as a temporary function. The return type, if specified, must be a
        struct type and it's validated when building the proto message
        for the PythonUDTF.
        """
        udtf = PythonUDTF(
            func=function,
            return_type=return_type,
            eval_type=eval_type,
            python_ver=get_python_ver(),
        )

        func = CommonInlineUserDefinedTableFunction(
            function_name=name,
            function=udtf,
            deterministic=deterministic,
            arguments=[],
        ).udtf_plan(self)

        req = self._execute_plan_request_with_metadata()
        req.plan.command.register_table_function.CopyFrom(func)

        self._execute(req)
        return name

    def register_data_source(self, dataSource: Type["DataSource"]) -> None:
        """
        Register a data source in the session catalog.
        """
        data_source = PythonDataSource(
            data_source=dataSource,
            python_ver=get_python_ver(),
        )
        proto = CommonInlineUserDefinedDataSource(
            name=dataSource.name(),
            data_source=data_source,
        ).to_data_source_proto(self)

        req = self._execute_plan_request_with_metadata()
        req.plan.command.register_data_source.CopyFrom(proto)
        self._execute(req)

    def register_java(
        self,
        name: str,
        javaClassName: str,
        return_type: Optional["DataTypeOrString"] = None,
        aggregate: bool = False,
    ) -> None:
        # construct a JavaUDF
        if return_type is None:
            java_udf = JavaUDF(class_name=javaClassName, aggregate=aggregate)
        else:
            java_udf = JavaUDF(class_name=javaClassName, output_type=return_type)
        fun = CommonInlineUserDefinedFunction(
            function_name=name,
            function=java_udf,
        ).to_plan_judf(self)
        # construct the request
        req = self._execute_plan_request_with_metadata()
        req.plan.command.register_function.CopyFrom(fun)

        self._execute(req)

    def _build_metrics(self, metrics: "pb2.ExecutePlanResponse.Metrics") -> Iterator[PlanMetrics]:
        return (
            PlanMetrics(
                x.name,
                x.plan_id,
                x.parent,
                [MetricValue(k, v.value, v.metric_type) for k, v in x.execution_metrics.items()],
            )
            for x in metrics.metrics
        )

    def _resources(self) -> Dict[str, ResourceInformation]:
        logger.debug("Fetching the resources")
        cmd = pb2.Command()
        cmd.get_resources_command.SetInParent()
        _, properties, _ = self.execute_command(cmd)
        resources = properties["get_resources_command_result"]
        return resources

    def to_table_as_iterator(
        self, plan: pb2.Plan, observations: Dict[str, Observation]
    ) -> Iterator[Union[StructType, "pa.Table"]]:
        """
        Return given plan as a PyArrow Table iterator.
        """
        if logger.isEnabledFor(logging.DEBUG):
            # inside an if statement to not incur a performance cost converting proto to string
            # when not at debug log level.
            logger.debug(f"Executing plan {self._proto_to_string(plan, True)}")
        req = self._execute_plan_request_with_metadata()
        req.plan.CopyFrom(plan)
        with Progress(handlers=self._progress_handlers, operation_id=req.operation_id) as progress:
            for response in self._execute_and_fetch_as_iterator(req, observations, progress):
                if isinstance(response, StructType):
                    yield response
                elif isinstance(response, pa.RecordBatch):
                    yield pa.Table.from_batches([response])

    def to_table(
        self, plan: pb2.Plan, observations: Dict[str, Observation]
    ) -> Tuple["pa.Table", Optional[StructType], ExecutionInfo]:
        """
        Return given plan as a PyArrow Table.
        """
        if logger.isEnabledFor(logging.DEBUG):
            # inside an if statement to not incur a performance cost converting proto to string
            # when not at debug log level.
            logger.debug(f"Executing plan {self._proto_to_string(plan, True)}")
        req = self._execute_plan_request_with_metadata()
        req.plan.CopyFrom(plan)
        table, schema, metrics, observed_metrics, _ = self._execute_and_fetch(req, observations)

        # Create a query execution object.
        ei = ExecutionInfo(metrics, observed_metrics)
        self._fire_execution_info(req.operation_id, ei)
        assert table is not None
        return table, schema, ei

    def to_pandas(
        self, plan: pb2.Plan, observations: Dict[str, Observation], **kwargs: Any
    ) -> Tuple["pd.DataFrame", "ExecutionInfo"]:
        """
        Return given plan as a pandas DataFrame.
        """
        if logger.isEnabledFor(logging.DEBUG):
            # inside an if statement to not incur a performance cost converting proto to string
            # when not at debug log level.
            logger.debug(f"Executing plan {self._proto_to_string(plan, True)}")
        req = self._execute_plan_request_with_metadata()
        req.plan.CopyFrom(plan)

        # Get all related configs in a batch
        (
            timezone,
            structHandlingMode,
            selfDestruct,
        ) = self.get_configs(
            "spark.sql.session.timeZone",
            "spark.sql.execution.pandas.structHandlingMode",
            "spark.sql.execution.arrow.pyspark.selfDestruct.enabled",
        )

        # if pandasStructHandlingMode is explicitly set, override the runtime config
        if "pandasStructHandlingMode" in kwargs:
            structHandlingMode = str(kwargs["pandasStructHandlingMode"])

        table, schema, metrics, observed_metrics, _ = self._execute_and_fetch(
            req, observations, selfDestruct == "true"
        )
        assert table is not None
        ei = ExecutionInfo(metrics, observed_metrics)
        self._fire_execution_info(req.operation_id, ei)

        schema = schema or from_arrow_schema(table.schema, prefer_timestamp_ntz=True)
        assert schema is not None and isinstance(schema, StructType)

        pdf = _convert_arrow_table_to_pandas(
            arrow_table=table,
            schema=schema,
            timezone=timezone,
            struct_handling_mode=structHandlingMode,
            date_as_object=False,
            self_destruct=selfDestruct == "true",
        )

        if len(metrics) > 0:
            pdf.attrs["metrics"] = metrics
        if len(observed_metrics) > 0:
            pdf.attrs["observed_metrics"] = observed_metrics
        return pdf, ei

    def _proto_to_string(self, p: google.protobuf.message.Message, truncate: bool = False) -> str:
        """
        Helper method to generate a one line string representation of the plan.

        Parameters
        ----------
        p : google.protobuf.message.Message
            Generic Message type
        truncate: bool
            Indicates whether to truncate the message

        Returns
        -------
        Single line string of the serialized proto message.
        """
        try:
            max_level = 8 if truncate else sys.maxsize
            p2 = self._truncate(p, max_level) if truncate else p
            return text_format.MessageToString(p2, as_one_line=True)
        except RecursionError:
            return "<Truncated message due to recursion error>"
        except Exception:
            return "<Truncated message due to truncation error>"

    def _truncate(
        self, p: google.protobuf.message.Message, allowed_recursion_depth: int
    ) -> google.protobuf.message.Message:
        """
        Helper method to truncate the protobuf message.
        Refer to 'org.apache.spark.sql.connect.common.Abbreviator' in the server side.
        """

        def truncate_str(s: str) -> str:
            if len(s) > 1024:
                return s[:1024] + "[truncated]"
            return s

        def truncate_bytes(b: bytes) -> bytes:
            if len(b) > 8:
                return b[:8] + b"[truncated]"
            return b

        p2 = copy.deepcopy(p)

        for descriptor, value in p.ListFields():
            if value is not None:
                field_name = descriptor.name

                if descriptor.type == descriptor.TYPE_MESSAGE:
                    if allowed_recursion_depth == 0:
                        p2.ClearField(field_name)
                    elif descriptor.label == descriptor.LABEL_REPEATED:
                        p2.ClearField(field_name)
                        getattr(p2, field_name).extend(
                            [self._truncate(v, allowed_recursion_depth - 1) for v in value]
                        )
                    else:
                        getattr(p2, field_name).CopyFrom(
                            self._truncate(value, allowed_recursion_depth - 1)
                        )

                elif descriptor.type == descriptor.TYPE_STRING:
                    if descriptor.label == descriptor.LABEL_REPEATED:
                        p2.ClearField(field_name)
                        getattr(p2, field_name).extend([truncate_str(v) for v in value])
                    else:
                        setattr(p2, field_name, truncate_str(value))

                elif descriptor.type == descriptor.TYPE_BYTES:
                    if descriptor.label == descriptor.LABEL_REPEATED:
                        p2.ClearField(field_name)
                        getattr(p2, field_name).extend([truncate_bytes(v) for v in value])
                    else:
                        setattr(p2, field_name, truncate_bytes(value))

        return p2

    def schema(self, plan: pb2.Plan) -> StructType:
        """
        Return schema for given plan.
        """
        if logger.isEnabledFor(logging.DEBUG):
            # inside an if statement to not incur a performance cost converting proto to string
            # when not at debug log level.
            logger.debug(f"Schema for plan: {self._proto_to_string(plan, True)}")
        schema = self._analyze(method="schema", plan=plan).schema
        assert schema is not None
        # Server side should populate the struct field which is the schema.
        assert isinstance(schema, StructType)
        return schema

    def explain_string(self, plan: pb2.Plan, explain_mode: str = "extended") -> str:
        """
        Return explain string for given plan.
        """
        if logger.isEnabledFor(logging.DEBUG):
            # inside an if statement to not incur a performance cost converting proto to string
            # when not at debug log level.
            logger.debug(
                f"Explain (mode={explain_mode}) for plan {self._proto_to_string(plan, True)}"
            )
        result = self._analyze(
            method="explain", plan=plan, explain_mode=explain_mode
        ).explain_string
        assert result is not None
        return result

    def execute_command(
        self, command: pb2.Command, observations: Optional[Dict[str, Observation]] = None
    ) -> Tuple[Optional[pd.DataFrame], Dict[str, Any], ExecutionInfo]:
        """
        Execute given command.
        """
        if logger.isEnabledFor(logging.DEBUG):
            # inside an if statement to not incur a performance cost converting proto to string
            # when not at debug log level.
            logger.debug(f"Execute command for command {self._proto_to_string(command, True)}")
        req = self._execute_plan_request_with_metadata()
        if self._user_id:
            req.user_context.user_id = self._user_id
        self._set_command_in_plan(req.plan, command)
        data, _, metrics, observed_metrics, properties = self._execute_and_fetch(
            req, observations or {}
        )
        # Create a query execution object.
        ei = ExecutionInfo(metrics, observed_metrics)
        self._fire_execution_info(req.operation_id, ei)
        if data is not None:
            return (data.to_pandas(), properties, ei)
        else:
            return (None, properties, ei)

    def execute_command_as_iterator(
        self, command: pb2.Command, observations: Optional[Dict[str, Observation]] = None
    ) -> Iterator[Dict[str, Any]]:
        """
        Execute given command. Similar to execute_command, but the value is returned using yield.
        """
        if logger.isEnabledFor(logging.DEBUG):
            # inside an if statement to not incur a performance cost converting proto to string
            # when not at debug log level.
            logger.debug(
                f"Execute command as iterator for command {self._proto_to_string(command, True)}"
            )
        req = self._execute_plan_request_with_metadata()
        if self._user_id:
            req.user_context.user_id = self._user_id
        self._set_command_in_plan(req.plan, command)
        for response in self._execute_and_fetch_as_iterator(req, observations or {}):
            if isinstance(response, dict):
                yield response
            else:
                raise PySparkValueError(
                    errorClass="UNKNOWN_RESPONSE",
                    messageParameters={
                        "response": str(response),
                    },
                )

    def same_semantics(self, plan: pb2.Plan, other: pb2.Plan) -> bool:
        """
        return if two plans have the same semantics.
        """
        result = self._analyze(method="same_semantics", plan=plan, other=other).is_same_semantics
        assert result is not None
        return result

    def semantic_hash(self, plan: pb2.Plan) -> int:
        """
        returns a `hashCode` of the logical query plan.
        """
        result = self._analyze(method="semantic_hash", plan=plan).semantic_hash
        assert result is not None
        return result

    def close(self) -> None:
        """
        Close the channel.
        """
        concurrent.futures.wait(self._release_futures, timeout=10)
        ExecutePlanResponseReattachableIterator.shutdown_threadpool_if_idle()
        self._channel.close()
        self._closed = True

    @property
    def is_closed(self) -> bool:
        """
        Returns if the channel was closed previously using close() method
        """
        return self._closed

    @property
    def host(self) -> str:
        """
        The hostname where this client intends to connect.
        """
        return self._builder.host

    @property
    def token(self) -> Optional[str]:
        """
        The authentication bearer token during connection.
        If authentication is not using a bearer token, None will be returned.
        """
        return self._builder.token

    def _update_request_with_user_context_extensions(
        self,
        req: Union[
            pb2.AnalyzePlanRequest,
            pb2.ConfigRequest,
            pb2.ExecutePlanRequest,
            pb2.FetchErrorDetailsRequest,
            pb2.InterruptRequest,
        ],
    ) -> None:
        with self.global_user_context_extensions_lock:
            for _, extension in self.global_user_context_extensions:
                req.user_context.extensions.append(extension)
        if not hasattr(self.thread_local, "user_context_extensions"):
            return
        for _, extension in self.thread_local.user_context_extensions:
            req.user_context.extensions.append(extension)

    def _execute_plan_request_with_metadata(
        self, operation_id: Optional[str] = None
    ) -> pb2.ExecutePlanRequest:
        req = pb2.ExecutePlanRequest(
            session_id=self._session_id,
            client_type=self._builder.userAgent,
            tags=list(self.get_tags()),
        )
        if self._server_session_id is not None:
            req.client_observed_server_side_session_id = self._server_session_id
        if self._user_id:
            req.user_context.user_id = self._user_id
        # Add request option to allow result chunking.
        req.request_options.append(
            pb2.ExecutePlanRequest.RequestOption(
                result_chunking_options=pb2.ResultChunkingOptions(
                    allow_arrow_batch_chunking=self._allow_arrow_batch_chunking,
                    preferred_arrow_chunk_size=self._preferred_arrow_chunk_size,
                )
            )
        )
        if operation_id is not None:
            try:
                uuid.UUID(operation_id, version=4)
            except ValueError as ve:
                raise PySparkValueError(
                    errorClass="INVALID_OPERATION_UUID_ID",
                    messageParameters={"arg_name": "operation_id", "origin": str(ve)},
                )
            req.operation_id = operation_id
        self._update_request_with_user_context_extensions(req)

        if call_stack_trace := self.__class__._build_call_stack_trace():
            req.user_context.extensions.append(call_stack_trace)
        return req

    def _analyze_plan_request_with_metadata(self) -> pb2.AnalyzePlanRequest:
        req = pb2.AnalyzePlanRequest()
        req.session_id = self._session_id
        if self._server_session_id is not None:
            req.client_observed_server_side_session_id = self._server_session_id
        req.client_type = self._builder.userAgent
        if self._user_id:
            req.user_context.user_id = self._user_id
        self._update_request_with_user_context_extensions(req)
        if call_stack_trace := self.__class__._build_call_stack_trace():
            req.user_context.extensions.append(call_stack_trace)
        return req

    def _analyze(self, method: str, **kwargs: Any) -> AnalyzeResult:
        """
        Call the analyze RPC of Spark Connect.

        Returns
        -------
        The result of the analyze call.
        """
        req = self._analyze_plan_request_with_metadata()
        if method == "schema":
            req.schema.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan")))
        elif method == "explain":
            req.explain.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan")))
            explain_mode = kwargs.get("explain_mode")
            if explain_mode not in ["simple", "extended", "codegen", "cost", "formatted"]:
                raise PySparkValueError(
                    errorClass="UNKNOWN_EXPLAIN_MODE",
                    messageParameters={
                        "explain_mode": str(explain_mode),
                    },
                )
            if explain_mode == "simple":
                req.explain.explain_mode = (
                    pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_SIMPLE
                )
            elif explain_mode == "extended":
                req.explain.explain_mode = (
                    pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_EXTENDED
                )
            elif explain_mode == "cost":
                req.explain.explain_mode = (
                    pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_COST
                )
            elif explain_mode == "codegen":
                req.explain.explain_mode = (
                    pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_CODEGEN
                )
            else:  # formatted
                req.explain.explain_mode = (
                    pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_FORMATTED
                )
        elif method == "tree_string":
            req.tree_string.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan")))
            level = kwargs.get("level")
            if level and isinstance(level, int):
                req.tree_string.level = level
        elif method == "is_local":
            req.is_local.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan")))
        elif method == "is_streaming":
            req.is_streaming.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan")))
        elif method == "input_files":
            req.input_files.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan")))
        elif method == "spark_version":
            req.spark_version.SetInParent()
        elif method == "ddl_parse":
            req.ddl_parse.ddl_string = cast(str, kwargs.get("ddl_string"))
        elif method == "same_semantics":
            req.same_semantics.target_plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan")))
            req.same_semantics.other_plan.CopyFrom(cast(pb2.Plan, kwargs.get("other")))
        elif method == "semantic_hash":
            req.semantic_hash.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan")))
        elif method == "persist":
            req.persist.relation.CopyFrom(cast(pb2.Relation, kwargs.get("relation")))
            if kwargs.get("storage_level", None) is not None:
                storage_level = cast(StorageLevel, kwargs.get("storage_level"))
                req.persist.storage_level.CopyFrom(storage_level_to_proto(storage_level))
        elif method == "unpersist":
            req.unpersist.relation.CopyFrom(cast(pb2.Relation, kwargs.get("relation")))
            if kwargs.get("blocking", None) is not None:
                req.unpersist.blocking = cast(bool, kwargs.get("blocking"))
        elif method == "get_storage_level":
            req.get_storage_level.relation.CopyFrom(cast(pb2.Relation, kwargs.get("relation")))
        elif method == "json_to_ddl":
            req.json_to_ddl.json_string = cast(str, kwargs.get("json_string"))
        else:
            raise PySparkValueError(
                errorClass="UNSUPPORTED_OPERATION",
                messageParameters={
                    "operation": method,
                },
            )

        try:
            for attempt in self._retrying():
                with attempt:
                    resp = self._stub.AnalyzePlan(req, metadata=self._builder.metadata())
                    self._verify_response_integrity(resp)
                    return AnalyzeResult.fromProto(resp)
            raise SparkConnectException("Invalid state during retry exception handling.")
        except Exception as error:
            self._handle_error(error)

    def _execute(self, req: pb2.ExecutePlanRequest) -> None:
        """
        Execute the passed request `req` and drop all results.

        Parameters
        ----------
        req : pb2.ExecutePlanRequest
            Proto representation of the plan.

        """
        logger.debug("Execute")

        for hook in self._session_hooks:
            req = hook.on_execute_plan(req)

        def handle_response(b: pb2.ExecutePlanResponse) -> None:
            self._verify_response_integrity(b)

        try:
            if self._use_reattachable_execute:
                # Don't use retryHandler - own retry handling is inside.
                generator = ExecutePlanResponseReattachableIterator(
                    req, self._stub, self._retrying, self._builder.metadata()
                )
                try:
                    for b in generator:
                        handle_response(b)
                finally:
                    generator.close()
                    self._release_futures.update(generator.release_futures)
            else:
                for attempt in self._retrying():
                    with attempt:
                        with disable_gc():
                            for b in self._stub.ExecutePlan(req, metadata=self._builder.metadata()):
                                handle_response(b)
        except Exception as error:
            self._handle_error(error)

    def _execute_and_fetch_as_iterator(
        self,
        req: pb2.ExecutePlanRequest,
        observations: Dict[str, Observation],
        progress: Optional["Progress"] = None,
    ) -> Iterator[
        Union[
            "pa.RecordBatch",
            StructType,
            PlanMetrics,
            PlanObservedMetrics,
            Dict[str, Any],
        ]
    ]:
        if logger.isEnabledFor(logging.DEBUG):
            # inside an if statement to not incur a performance cost converting proto to string
            # when not at debug log level.
            logger.debug(f"ExecuteAndFetchAsIterator. Request: {self._proto_to_string(req)}")

        for hook in self._session_hooks:
            req = hook.on_execute_plan(req)

        num_records = 0
        arrow_batch_chunks_to_assemble: List[bytes] = []

        def handle_response(
            b: pb2.ExecutePlanResponse,
        ) -> Iterator[
            Union[
                "pa.RecordBatch",
                StructType,
                PlanMetrics,
                PlanObservedMetrics,
                Dict[str, Any],
                any_pb2.Any,
            ]
        ]:
            nonlocal num_records
            # The session ID is the local session ID and should match what we expect.
            self._verify_response_integrity(b)
            if logger.isEnabledFor(logging.DEBUG):
                # inside an if statement to not incur a performance cost converting proto to string
                # when not at debug log level.
                logger.debug(
                    f"ExecuteAndFetchAsIterator. Response received: {self._proto_to_string(b)}"
                )

            if b.HasField("metrics"):
                logger.debug("Received metric batch.")
                yield from self._build_metrics(b.metrics)
            if b.observed_metrics:
                logger.debug("Received observed metric batch.")
                for x in b.observed_metrics:
                    observed_metrics = PlanObservedMetrics(
                        x.name, [v for v in x.values], list(x.keys)
                    )
                    if x.HasField("root_error_idx"):
                        if x.name in observations:
                            converted = convert_observation_errors(x.root_error_idx, list(x.errors))
                            observations[x.name]._set_error(converted)
                    else:
                        if observed_metrics.name == "__python_accumulator__":
                            for metric in observed_metrics.metrics:
                                aid, update = pickleSer.loads(LiteralExpression._to_value(metric))
                                if aid == SpecialAccumulatorIds.SQL_UDF_PROFIER_V2:
                                    self._profiler_collector._update(update)
                        elif observed_metrics.name in observations:
                            observation_result = observations[observed_metrics.name]._result
                            assert observation_result is not None
                            observation_result.update(
                                {
                                    key: LiteralExpression._to_value(metric)
                                    for key, metric in zip(
                                        observed_metrics.keys,
                                        observed_metrics.metrics,
                                    )
                                }
                            )
                    yield observed_metrics
            if b.HasField("schema"):
                logger.debug("Received the schema.")
                dt = types.proto_schema_to_pyspark_data_type(b.schema)
                assert isinstance(dt, StructType)
                yield dt
            if b.HasField("sql_command_result"):
                logger.debug("Received the SQL command result.")
                yield {"sql_command_result": b.sql_command_result.relation}
            if b.HasField("write_stream_operation_start_result"):
                field = "write_stream_operation_start_result"
                yield {field: b.write_stream_operation_start_result}
            if b.HasField("streaming_query_command_result"):
                yield {"streaming_query_command_result": b.streaming_query_command_result}
            if b.HasField("streaming_query_manager_command_result"):
                cmd_result = b.streaming_query_manager_command_result
                yield {"streaming_query_manager_command_result": cmd_result}
            if b.HasField("streaming_query_listener_events_result"):
                event_result = b.streaming_query_listener_events_result
                yield {"streaming_query_listener_events_result": event_result}
            if b.HasField("pipeline_command_result"):
                yield {"pipeline_command_result": b.pipeline_command_result}
            if b.HasField("pipeline_event_result"):
                yield {"pipeline_event_result": b.pipeline_event_result}
            if b.HasField("get_resources_command_result"):
                resources = {}
                for key, resource in b.get_resources_command_result.resources.items():
                    name = resource.name
                    addresses = [address for address in resource.addresses]
                    resources[key] = ResourceInformation(name, addresses)
                yield {"get_resources_command_result": resources}
            if b.HasField("extension"):
                yield b.extension
            if b.HasField("execution_progress"):
                if progress:
                    p = from_proto(b.execution_progress)
                    progress.update_ticks(*p, operation_id=b.operation_id)
            if b.HasField("arrow_batch"):
                logger.debug(
                    f"Received arrow batch rows={b.arrow_batch.row_count} "
                    f"Number of chunks in batch={b.arrow_batch.num_chunks_in_batch} "
                    f"Chunk index={b.arrow_batch.chunk_index} "
                    f"size={len(b.arrow_batch.data)}"
                )

                if arrow_batch_chunks_to_assemble:
                    # Expect next chunk of the same batch
                    if b.arrow_batch.chunk_index != len(arrow_batch_chunks_to_assemble):
                        raise SparkConnectException(
                            f"Expected chunk index {len(arrow_batch_chunks_to_assemble)} of the "
                            f"arrow batch but got {b.arrow_batch.chunk_index}."
                        )
                else:
                    # Expect next batch
                    if (
                        b.arrow_batch.HasField("start_offset")
                        and num_records != b.arrow_batch.start_offset
                    ):
                        # Expect next batch
                        raise SparkConnectException(
                            f"Expected arrow batch to start at row offset {num_records} in "
                            + "results, but received arrow batch starting at offset "
                            + f"{b.arrow_batch.start_offset}."
                        )
                    if b.arrow_batch.chunk_index != 0:
                        raise SparkConnectException(
                            f"Expected chunk index 0 of the next arrow batch "
                            f"but got {b.arrow_batch.chunk_index}."
                        )

                arrow_batch_chunks_to_assemble.append(b.arrow_batch.data)
                # Assemble the chunks to an arrow batch to process if
                # (a) chunking is not enabled (num_chunks_in_batch is not set or is 0,
                #  in this case, it is the single chunk in the batch)
                # (b) or the client has received all chunks of the batch.
                if (
                    not b.arrow_batch.HasField("num_chunks_in_batch")
                    or b.arrow_batch.num_chunks_in_batch == 0
                    or len(arrow_batch_chunks_to_assemble) == b.arrow_batch.num_chunks_in_batch
                ):
                    arrow_batch_data = b"".join(arrow_batch_chunks_to_assemble)
                    arrow_batch_chunks_to_assemble.clear()
                    logger.debug(
                        f"Assembling arrow batch of size {len(arrow_batch_data)} from "
                        f"{b.arrow_batch.num_chunks_in_batch} chunks."
                    )

                    num_records_in_batch = 0
                    with pa.ipc.open_stream(arrow_batch_data) as reader:
                        for batch in reader:
                            assert isinstance(batch, pa.RecordBatch)
                            num_records_in_batch += batch.num_rows
                            if num_records_in_batch != b.arrow_batch.row_count:
                                raise SparkConnectException(
                                    f"Expected {b.arrow_batch.row_count} rows in arrow batch but "
                                    + f"got {num_records_in_batch}."
                                )
                            num_records += num_records_in_batch
                            yield batch
            if b.HasField("create_resource_profile_command_result"):
                profile_id = b.create_resource_profile_command_result.profile_id
                yield {"create_resource_profile_command_result": profile_id}
            if b.HasField("checkpoint_command_result"):
                yield {
                    "checkpoint_command_result": proto_to_remote_cached_dataframe(
                        b.checkpoint_command_result.relation
                    )
                }
            if b.HasField("ml_command_result"):
                yield {"ml_command_result": b.ml_command_result}

        try:
            if self._use_reattachable_execute:
                # Don't use retryHandler - own retry handling is inside.
                generator = ExecutePlanResponseReattachableIterator(
                    req, self._stub, self._retrying, self._builder.metadata()
                )
                try:
                    for b in generator:
                        yield from handle_response(b)
                finally:
                    generator.close()
                    self._release_futures.update(generator.release_futures)
            else:
                for attempt in self._retrying():
                    with attempt:
                        with disable_gc():
                            it = iter(
                                self._stub.ExecutePlan(req, metadata=self._builder.metadata())
                            )
                        while True:
                            try:
                                with disable_gc():
                                    b = next(it)
                                yield from handle_response(b)
                            except StopIteration:
                                break
        except KeyboardInterrupt as kb:
            logger.debug(f"Interrupt request received for operation={req.operation_id}")
            if progress is not None:
                progress.finish()
            self.interrupt_operation(req.operation_id)
            raise kb
        except Exception as error:
            self._handle_error(error)

    def _execute_and_fetch(
        self,
        req: pb2.ExecutePlanRequest,
        observations: Dict[str, Observation],
        self_destruct: bool = False,
    ) -> Tuple[
        Optional["pa.Table"],
        Optional[StructType],
        List[PlanMetrics],
        List[PlanObservedMetrics],
        Dict[str, Any],
    ]:
        logger.debug("ExecuteAndFetch")

        observed_metrics: List[PlanObservedMetrics] = []
        metrics: List[PlanMetrics] = []
        batches: List[pa.RecordBatch] = []
        schema: Optional[StructType] = None
        properties: Dict[str, Any] = {}

        with Progress(handlers=self._progress_handlers, operation_id=req.operation_id) as progress:
            for response in self._execute_and_fetch_as_iterator(
                req, observations, progress=progress
            ):
                if isinstance(response, StructType):
                    schema = response
                elif isinstance(response, pa.RecordBatch):
                    batches.append(response)
                elif isinstance(response, PlanMetrics):
                    metrics.append(response)
                elif isinstance(response, PlanObservedMetrics):
                    observed_metrics.append(response)
                elif isinstance(response, dict):
                    properties.update(**response)
                else:
                    raise PySparkValueError(
                        errorClass="UNKNOWN_RESPONSE",
                        messageParameters={
                            "response": response,
                        },
                    )

        if len(batches) > 0:
            if self_destruct:
                results = []
                for batch in batches:
                    # self_destruct frees memory column-wise, but Arrow record batches are
                    # oriented row-wise, so copies each column into its own allocation
                    batch = pa.RecordBatch.from_arrays(
                        [
                            # This call actually reallocates the array
                            pa.concat_arrays([array])
                            for array in batch
                        ],
                        schema=batch.schema,
                    )
                    results.append(batch)
                table = pa.Table.from_batches(batches=results)
                # Ensure only the table has a reference to the batches, so that
                # self_destruct (if enabled) is effective
                del results
                del batches
            else:
                table = pa.Table.from_batches(batches=batches)
            return table, schema, metrics, observed_metrics, properties
        else:
            return None, schema, metrics, observed_metrics, properties

    def _config_request_with_metadata(self) -> pb2.ConfigRequest:
        req = pb2.ConfigRequest()
        req.session_id = self._session_id
        if self._server_session_id is not None:
            req.client_observed_server_side_session_id = self._server_session_id
        req.client_type = self._builder.userAgent
        if self._user_id:
            req.user_context.user_id = self._user_id
        self._update_request_with_user_context_extensions(req)
        if call_stack_trace := self.__class__._build_call_stack_trace():
            req.user_context.extensions.append(call_stack_trace)
        return req

    def get_configs(self, *keys: str) -> Tuple[Optional[str], ...]:
        op = pb2.ConfigRequest.Operation(get=pb2.ConfigRequest.Get(keys=keys))
        configs = dict(self.config(op).pairs)
        return tuple(configs.get(key) for key in keys)

    def get_config_dict(self, *keys: str) -> Mapping[str, Optional[str]]:
        op = pb2.ConfigRequest.Operation(get=pb2.ConfigRequest.Get(keys=keys))
        return dict(self.config(op).pairs)

    def get_config_with_defaults(
        self, *pairs: Tuple[str, Optional[str]]
    ) -> Tuple[Optional[str], ...]:
        op = pb2.ConfigRequest.Operation(
            get_with_default=pb2.ConfigRequest.GetWithDefault(
                pairs=[pb2.KeyValue(key=key, value=default) for key, default in pairs]
            )
        )
        configs = dict(self.config(op).pairs)
        return tuple(configs.get(key) for key, _ in pairs)

    def config(self, operation: pb2.ConfigRequest.Operation) -> ConfigResult:
        """
        Call the config RPC of Spark Connect.

        Parameters
        ----------
        operation : str
           Operation kind

        Returns
        -------
        The result of the config call.
        """
        req = self._config_request_with_metadata()
        if self._server_session_id is not None:
            req.client_observed_server_side_session_id = self._server_session_id
        req.operation.CopyFrom(operation)
        try:
            for attempt in self._retrying():
                with attempt:
                    resp = self._stub.Config(req, metadata=self._builder.metadata())
                    self._verify_response_integrity(resp)
                    return ConfigResult.fromProto(resp)
            raise SparkConnectException("Invalid state during retry exception handling.")
        except Exception as error:
            self._handle_error(error)

    def _interrupt_request(
        self, interrupt_type: str, id_or_tag: Optional[str] = None
    ) -> pb2.InterruptRequest:
        req = pb2.InterruptRequest()
        req.session_id = self._session_id
        if self._server_session_id is not None:
            req.client_observed_server_side_session_id = self._server_session_id
        req.client_type = self._builder.userAgent
        if interrupt_type == "all":
            req.interrupt_type = pb2.InterruptRequest.InterruptType.INTERRUPT_TYPE_ALL
        elif interrupt_type == "tag":
            assert id_or_tag is not None
            req.interrupt_type = pb2.InterruptRequest.InterruptType.INTERRUPT_TYPE_TAG
            req.operation_tag = id_or_tag
        elif interrupt_type == "operation":
            assert id_or_tag is not None
            req.interrupt_type = pb2.InterruptRequest.InterruptType.INTERRUPT_TYPE_OPERATION_ID
            req.operation_id = id_or_tag
        else:
            raise PySparkValueError(
                errorClass="UNKNOWN_INTERRUPT_TYPE",
                messageParameters={
                    "interrupt_type": str(interrupt_type),
                },
            )
        if self._user_id:
            req.user_context.user_id = self._user_id
        self._update_request_with_user_context_extensions(req)
        return req

    def interrupt_all(self) -> Optional[List[str]]:
        req = self._interrupt_request("all")
        try:
            for attempt in self._retrying():
                with attempt:
                    resp = self._stub.Interrupt(req, metadata=self._builder.metadata())
                    self._verify_response_integrity(resp)
                    return list(resp.interrupted_ids)
            raise SparkConnectException("Invalid state during retry exception handling.")
        except Exception as error:
            self._handle_error(error)

    def interrupt_tag(self, tag: str) -> Optional[List[str]]:
        req = self._interrupt_request("tag", tag)
        try:
            for attempt in self._retrying():
                with attempt:
                    resp = self._stub.Interrupt(req, metadata=self._builder.metadata())
                    self._verify_response_integrity(resp)
                    return list(resp.interrupted_ids)
            raise SparkConnectException("Invalid state during retry exception handling.")
        except Exception as error:
            self._handle_error(error)

    def interrupt_operation(self, op_id: str) -> Optional[List[str]]:
        req = self._interrupt_request("operation", op_id)
        try:
            for attempt in self._retrying():
                with attempt:
                    resp = self._stub.Interrupt(req, metadata=self._builder.metadata())
                    self._verify_response_integrity(resp)
                    return list(resp.interrupted_ids)
            raise SparkConnectException("Invalid state during retry exception handling.")
        except Exception as error:
            self._handle_error(error)

    def release_session(self) -> None:
        req = pb2.ReleaseSessionRequest()
        req.session_id = self._session_id
        req.client_type = self._builder.userAgent
        if self._user_id:
            req.user_context.user_id = self._user_id
        try:
            for attempt in self._retrying():
                with attempt:
                    resp = self._stub.ReleaseSession(req, metadata=self._builder.metadata())
                    self._verify_response_integrity(resp)
                    return
            raise SparkConnectException("Invalid state during retry exception handling.")
        except Exception as error:
            self._handle_error(error)

    def _get_operation_statuses(
        self,
        operation_ids: Optional[List[str]] = None,
        operation_extensions: Optional[List[any_pb2.Any]] = None,
        request_extensions: Optional[List[any_pb2.Any]] = None,
    ) -> "pb2.GetStatusResponse":
        """
        Get status of operations in the session.

        Parameters
        ----------
        operation_ids : list of str, optional
            List of operation IDs to get status for.
            If None or empty, returns status of all operations in the session.
        operation_extensions : list of google.protobuf.any_pb2.Any, optional
            Per-operation extension messages to include in the OperationStatusRequest to request
            additional per-operation information.
        request_extensions : list of google.protobuf.any_pb2.Any, optional
            Request-level extension messages to include in the GetStatusRequest.

        Returns
        -------
        pb2.GetStatusResponse
            The full GetStatusResponse, including operation_statuses and any extensions.
        """
        req = pb2.GetStatusRequest()
        req.session_id = self._session_id
        req.client_type = self._builder.userAgent
        if self._user_id:
            req.user_context.user_id = self._user_id
        if self._server_session_id:
            req.client_observed_server_side_session_id = self._server_session_id

        req.operation_status.SetInParent()

        if operation_ids:
            req.operation_status.operation_ids.extend(operation_ids)
        if operation_extensions:
            req.operation_status.extensions.extend(operation_extensions)
        if request_extensions:
            req.extensions.extend(request_extensions)

        try:
            for attempt in self._retrying():
                with attempt:
                    resp = self._stub.GetStatus(req, metadata=self._builder.metadata())
                    self._verify_response_integrity(resp)
                    return resp
            raise SparkConnectException("Invalid state during retry exception handling.")
        except Exception as error:
            self._handle_error(error)

    def add_tag(self, tag: str) -> None:
        self._throw_if_invalid_tag(tag)
        if not hasattr(self.thread_local, "tags"):
            self.thread_local.tags = set()
        self.thread_local.tags.add(tag)

    def remove_tag(self, tag: str) -> None:
        self._throw_if_invalid_tag(tag)
        if not hasattr(self.thread_local, "tags"):
            self.thread_local.tags = set()
        self.thread_local.tags.remove(tag)

    def get_tags(self) -> Set[str]:
        if not hasattr(self.thread_local, "tags"):
            self.thread_local.tags = set()
        return self.thread_local.tags

    def clear_tags(self) -> None:
        self.thread_local.tags = set()

    def _throw_if_invalid_tag(self, tag: str) -> None:
        """
        Validate if a tag for ExecutePlanRequest.tags is valid. Throw ``ValueError`` if
        not.
        """
        spark_job_tags_sep = ","
        if tag is None:
            raise PySparkValueError(
                errorClass="CANNOT_BE_NONE", message_paramters={"arg_name": "Spark Connect tag"}
            )
        if spark_job_tags_sep in tag:
            raise PySparkValueError(
                errorClass="VALUE_ALLOWED",
                messageParameters={
                    "arg_name": "Spark Connect tag",
                    "disallowed_value": spark_job_tags_sep,
                },
            )
        if len(tag) == 0:
            raise PySparkValueError(
                errorClass="VALUE_NOT_NON_EMPTY_STR",
                messageParameters={"arg_name": "Spark Connect tag", "arg_value": tag},
            )

    def add_threadlocal_user_context_extension(self, extension: any_pb2.Any) -> str:
        if not hasattr(self.thread_local, "user_context_extensions"):
            self.thread_local.user_context_extensions = list()
        extension_id = "threadlocal_" + str(uuid.uuid4())
        self.thread_local.user_context_extensions.append((extension_id, extension))
        return extension_id

    def add_global_user_context_extension(self, extension: any_pb2.Any) -> str:
        extension_id = "global_" + str(uuid.uuid4())
        with self.global_user_context_extensions_lock:
            self.global_user_context_extensions.append((extension_id, extension))
        return extension_id

    def remove_user_context_extension(self, extension_id: str) -> None:
        if extension_id.find("threadlocal_") == 0:
            if not hasattr(self.thread_local, "user_context_extensions"):
                return
            self.thread_local.user_context_extensions = list(
                filter(lambda ex: ex[0] != extension_id, self.thread_local.user_context_extensions)
            )
        elif extension_id.find("global_") == 0:
            with self.global_user_context_extensions_lock:
                self.global_user_context_extensions = list(
                    filter(lambda ex: ex[0] != extension_id, self.global_user_context_extensions)
                )

    def clear_user_context_extensions(self) -> None:
        if hasattr(self.thread_local, "user_context_extensions"):
            self.thread_local.user_context_extensions = list()
        with self.global_user_context_extensions_lock:
            self.global_user_context_extensions = list()

    def _handle_error(self, error: Exception) -> NoReturn:
        """
        Handle errors that occur during RPC calls.

        Parameters
        ----------
        error : Exception
            An exception thrown during RPC calls.

        Returns
        -------
        Throws the appropriate internal Python exception.
        """

        if getattr(self.thread_local, "inside_error_handling", False):
            # We are already inside error handling routine,
            # avoid recursive error processing (with potentially infinite recursion)
            raise error

        try:
            self.thread_local.inside_error_handling = True
            if isinstance(error, grpc.RpcError):
                self._handle_rpc_error(error)
            raise error
        finally:
            self.thread_local.inside_error_handling = False

    def _fetch_enriched_error(self, info: "ErrorInfo") -> Optional[pb2.FetchErrorDetailsResponse]:
        if "errorId" not in info.metadata:
            return None

        req = pb2.FetchErrorDetailsRequest(
            session_id=self._session_id,
            client_type=self._builder.userAgent,
            error_id=info.metadata["errorId"],
        )
        if self._server_session_id is not None:
            req.client_observed_server_side_session_id = self._server_session_id
        if self._user_id:
            req.user_context.user_id = self._user_id
        self._update_request_with_user_context_extensions(req)
        try:
            return self._stub.FetchErrorDetails(req, metadata=self._builder.metadata())
        except grpc.RpcError:
            return None

    def _display_server_stack_trace(self) -> bool:
        from pyspark.sql.connect.conf import RuntimeConf

        conf = RuntimeConf(self)
        try:
            if conf.get("spark.sql.connect.serverStacktrace.enabled") == "true":
                return True
            return conf.get("spark.sql.pyspark.jvmStacktrace.enabled") == "true"
        except Exception as e:  # noqa: F841
            # Falls back to true if an exception occurs during reading the config.
            # Otherwise, it will recursively try to get the conf when it consistently
            # fails, ending up with `RecursionError`.
            return True

    def _handle_rpc_error(self, rpc_error: grpc.RpcError) -> NoReturn:
        """
        Error handling helper for dealing with GRPC Errors. On the server side, certain
        exceptions are enriched with additional RPC Status information. These are
        unpacked in this function and put into the exception.

        To avoid overloading the user with GRPC errors, this message explicitly
        swallows the error context from the call. This GRPC Error is logged however,
        and can be enabled.

        Parameters
        ----------
        rpc_error : grpc.RpcError
           RPC Error containing the details of the exception.

        Returns
        -------
        Throws the appropriate internal Python exception.
        """
        logger.exception("GRPC Error received")
        # We have to cast the value here because, a RpcError is a Call as well.
        # https://grpc.github.io/grpc/python/grpc.html#grpc.UnaryUnaryMultiCallable.__call__
        error: grpc.Call = cast(grpc.Call, rpc_error)
        status_code: grpc.StatusCode = error.code()
        status: Optional[Status] = rpc_status.from_call(error)
        if status:
            for d in status.details:
                if d.Is(error_details_pb2.ErrorInfo.DESCRIPTOR):
                    info = error_details_pb2.ErrorInfo()
                    d.Unpack(info)
                    logger.debug(f"Received ErrorInfo: {info}")

                    if info.metadata.get("errorClass") == "INVALID_HANDLE.SESSION_CHANGED":
                        self._closed = True

                    if info.metadata.get("errorClass") == "CONNECT_INVALID_PLAN.CANNOT_PARSE":
                        # Disable plan compression if the server fails to interpret the plan.
                        logger.info(
                            "Disabling plan compression for the session due to "
                            "CONNECT_INVALID_PLAN.CANNOT_PARSE error."
                        )
                        self._plan_compression_threshold, self._plan_compression_algorithm = (
                            -1,
                            "NONE",
                        )

                    raise convert_exception(
                        info,
                        status.message,
                        self._fetch_enriched_error(info),
                        self._display_server_stack_trace(),
                        status_code,
                    ) from None

            raise SparkConnectGrpcException(
                message=status.message,
                grpc_status_code=status_code,
            ) from None
        else:
            raise SparkConnectGrpcException(
                message=str(error),
                grpc_status_code=status_code,
            ) from None

    def add_artifacts(self, *paths: str, pyfile: bool, archive: bool, file: bool) -> None:
        try:
            for path in paths:
                for attempt in self._retrying():
                    with attempt:
                        self._artifact_manager.add_artifacts(
                            path, pyfile=pyfile, archive=archive, file=file
                        )
        except Exception as error:
            self._handle_error(error)

    def copy_from_local_to_fs(self, local_path: str, dest_path: str) -> None:
        for attempt in self._retrying():
            with attempt:
                self._artifact_manager._add_forward_to_fs_artifacts(local_path, dest_path)

    def cache_artifact(self, blob: bytes) -> str:
        for attempt in self._retrying():
            with attempt:
                return self._artifact_manager.cache_artifact(blob)
        raise SparkConnectException("Invalid state during retry exception handling.")

    def cache_artifacts(self, blobs: list[bytes]) -> list[str]:
        for attempt in self._retrying():
            with attempt:
                return self._artifact_manager.cache_artifacts(blobs)
        raise SparkConnectException("Invalid state during retry exception handling.")

    def _verify_response_integrity(
        self,
        response: Union[
            pb2.ConfigResponse,
            pb2.ExecutePlanResponse,
            pb2.InterruptResponse,
            pb2.ReleaseExecuteResponse,
            pb2.AddArtifactsResponse,
            pb2.AnalyzePlanResponse,
            pb2.FetchErrorDetailsResponse,
            pb2.ReleaseSessionResponse,
            pb2.GetStatusResponse,
        ],
    ) -> None:
        """
        Verifies the integrity of the response. This method checks if the session ID and the
        server-side session ID match. If not, it throws an exception.
        Parameters
        ----------
        response - One of the different response types handled by the Spark Connect service
        """
        if self._session_id != response.session_id:
            raise PySparkAssertionError(
                "Received incorrect session identifier for request:"
                f"{response.session_id} != {self._session_id}"
            )
        if self._server_session_id is not None:
            if (
                response.server_side_session_id
                and response.server_side_session_id != self._server_session_id
            ):
                self._closed = True
                raise PySparkAssertionError(
                    "Received incorrect server side session identifier for request. "
                    "Please create a new Spark Session to reconnect. ("
                    f"{response.server_side_session_id} != {self._server_session_id})"
                )
        else:
            # Update the server side session ID.
            self._server_session_id = response.server_side_session_id

    def _create_profile(self, profile: pb2.ResourceProfile) -> int:
        """Create the ResourceProfile on the server side and return the profile ID"""
        logger.debug("Creating the ResourceProfile")
        cmd = pb2.Command()
        cmd.create_resource_profile_command.profile.CopyFrom(profile)
        _, properties, _ = self.execute_command(cmd)
        profile_id = properties["create_resource_profile_command_result"]
        return profile_id

    def _delete_ml_cache(self, cache_ids: List[str], evict_only: bool = False) -> List[str]:
        # try best to delete the cache
        try:
            if len(cache_ids) > 0:
                command = pb2.Command()
                command.ml_command.delete.obj_refs.extend(
                    [pb2.ObjectRef(id=cache_id) for cache_id in cache_ids]
                )
                command.ml_command.delete.evict_only = evict_only
                _, properties, _ = self.execute_command(command)

                assert properties is not None

                if properties is not None and "ml_command_result" in properties:
                    ml_command_result = properties["ml_command_result"]
                    deleted = ml_command_result.operator_info.obj_ref.id.split(",")
                    return cast(List[str], deleted)
            return []
        except Exception:
            return []

    def _on_exit(self) -> None:
        self._cleanup_ml_cache()
        if self._release_session_on_exit and not self._closed:
            try:
                self.release_session()
            except Exception:
                pass
            try:
                self.close()
            except Exception:
                pass

    def _cleanup_ml_cache(self) -> None:
        try:
            command = pb2.Command()
            command.ml_command.clean_cache.SetInParent()
            self.execute_command(command)
        except Exception:
            pass

    def _get_ml_cache_info(self) -> List[str]:
        command = pb2.Command()
        command.ml_command.get_cache_info.SetInParent()
        _, properties, _ = self.execute_command(command)

        assert properties is not None

        if properties is not None and "ml_command_result" in properties:
            ml_command_result = properties["ml_command_result"]
            return [item.string for item in ml_command_result.param.array.elements]

        return []

    def _query_model_size(self, model_ref_id: str) -> int:
        command = pb2.Command()
        command.ml_command.get_model_size.CopyFrom(
            pb2.MlCommand.GetModelSize(model_ref=pb2.ObjectRef(id=model_ref_id))
        )
        _, properties, _ = self.execute_command(command)

        assert properties is not None

        ml_command_result = properties["ml_command_result"]
        return ml_command_result.param.long

    def _set_relation_in_plan(self, plan: pb2.Plan, relation: pb2.Relation) -> None:
        """Sets the relation in the plan, attempting compression if configured."""
        self._try_compress_and_set_plan(
            plan=plan,
            message=relation,
            op_type=pb2.Plan.CompressedOperation.OpType.OP_TYPE_RELATION,
        )

    def _set_command_in_plan(self, plan: pb2.Plan, command: pb2.Command) -> None:
        """Sets the command in the plan, attempting compression if configured."""
        self._try_compress_and_set_plan(
            plan=plan,
            message=command,
            op_type=pb2.Plan.CompressedOperation.OpType.OP_TYPE_COMMAND,
        )

    def _try_compress_and_set_plan(
        self,
        plan: pb2.Plan,
        message: google.protobuf.message.Message,
        op_type: pb2.Plan.CompressedOperation.OpType.ValueType,
    ) -> None:
        """
        Tries to compress a protobuf message and sets it on the plan.
        If compression is not enabled, not effective, or not available,
        it falls back to the original message.
        """
        (
            plan_compression_threshold,
            plan_compression_algorithm,
        ) = self._get_plan_compression_threshold_and_algorithm()
        plan_compression_enabled = (
            plan_compression_threshold is not None
            and plan_compression_threshold >= 0
            and plan_compression_algorithm is not None
            and plan_compression_algorithm != "NONE"
        )
        if plan_compression_enabled:
            serialized_msg = message.SerializeToString()
            original_size = len(serialized_msg)
            if (
                original_size > plan_compression_threshold
                and plan_compression_algorithm == "ZSTD"
                and self._zstd_module
            ):
                start_time = time.time()
                compressed_operation = pb2.Plan.CompressedOperation(
                    data=self._zstd_module.compress(serialized_msg),
                    op_type=op_type,
                    compression_codec=pb2.CompressionCodec.COMPRESSION_CODEC_ZSTD,
                )
                duration = time.time() - start_time
                compressed_size = len(compressed_operation.data)
                logger.debug(
                    f"Plan compression: original_size={original_size}, "
                    f"compressed_size={compressed_size}, "
                    f"saving_ratio={1 - compressed_size / original_size:.2f}, "
                    f"duration_s={duration:.1f}"
                )
                if compressed_size < original_size:
                    plan.compressed_operation.CopyFrom(compressed_operation)
                    return
                else:
                    logger.debug("Plan compression not effective. Using original plan.")

        if op_type == pb2.Plan.CompressedOperation.OpType.OP_TYPE_RELATION:
            plan.root.CopyFrom(message)  # type: ignore[arg-type]
        else:
            plan.command.CopyFrom(message)  # type: ignore[arg-type]

    def _get_plan_compression_threshold_and_algorithm(self) -> Tuple[int, str]:
        if self._plan_compression_threshold is None or self._plan_compression_algorithm is None:
            try:
                (
                    plan_compression_threshold_str,
                    self._plan_compression_algorithm,
                ) = self.get_configs(
                    "spark.connect.session.planCompression.threshold",
                    "spark.connect.session.planCompression.defaultAlgorithm",
                )
                self._plan_compression_threshold = (
                    int(plan_compression_threshold_str) if plan_compression_threshold_str else -1
                )
                logger.debug(
                    f"Plan compression threshold: {self._plan_compression_threshold}, "
                    f"algorithm: {self._plan_compression_algorithm}"
                )
            except Exception as e:
                self._plan_compression_threshold = -1
                self._plan_compression_algorithm = "NONE"
                logger.debug(
                    "Plan compression is disabled because the server does not support it.", e
                )
        return (
            self._plan_compression_threshold,
            self._plan_compression_algorithm,
        )  # type: ignore[return-value]

    def clone(self, new_session_id: Optional[str] = None) -> "SparkConnectClient":
        """
        Clone this client session on the server side. The server-side session is cloned with
        all its current state (SQL configurations, temporary views, registered functions,
        catalog state) copied over to a new independent session. The returned client with the
        cloned session is isolated from this client's session - any subsequent changes to
        either session's server-side state will not be reflected in the other.

        Parameters
        ----------
        new_session_id : str, optional
            Custom session ID to use for the cloned session (must be a valid UUID).
            If not provided, a new UUID will be generated.

        Returns
        -------
        SparkConnectClient
            A new SparkConnectClient instance with the cloned session.

        Notes
        -----
        This creates a new server-side session with the specified or generated session ID
        while preserving the current session's configuration and state.

        .. note::
            This is a developer API.
        """
        from pyspark.sql.connect.proto import base_pb2 as pb2

        request = pb2.CloneSessionRequest(
            session_id=self._session_id,
            client_type="python",
        )
        if self._user_id is not None:
            request.user_context.user_id = self._user_id

        if new_session_id is not None:
            request.new_session_id = new_session_id

        for attempt in self._retrying():
            with attempt:
                response: pb2.CloneSessionResponse = self._stub.CloneSession(
                    request, metadata=self._builder.metadata()
                )

        # Assert that the returned session ID matches the requested ID if one was provided
        if new_session_id is not None:
            assert response.new_session_id == new_session_id, (
                f"Returned session ID '{response.new_session_id}' does not match "
                f"requested ID '{new_session_id}'"
            )

        # Create a new client with the cloned session ID
        new_connection = copy.deepcopy(self._builder)
        new_connection.set(ChannelBuilder.PARAM_SESSION_ID, response.new_session_id)

        # Create new client and explicitly set the session ID
        new_client = SparkConnectClient(
            connection=new_connection,
            user_id=self._user_id,
            use_reattachable_execute=self._use_reattachable_execute,
        )
        # Ensure the session ID is correctly set from the response
        new_client._session_id = response.new_session_id
        return new_client