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

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

import os
import unittest
import uuid
from typing import Optional

from pyspark.util import is_remote_only
from pyspark.errors import PySparkException
from pyspark.sql import SparkSession as PySparkSession
from pyspark.testing.connectutils import (
    should_test_connect,
    ReusedConnectTestCase,
    connect_requirement_message,
)
from pyspark.testing.utils import timeout

if should_test_connect:
    import grpc
    from pyspark.sql.connect.session import SparkSession as RemoteSparkSession
    from pyspark.sql.connect.client import ChannelBuilder, DefaultChannelBuilder
    from pyspark.errors.exceptions.connect import (
        AnalysisException,
        SparkConnectException,
        SparkConnectGrpcException,
        SparkUpgradeException,
    )

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


@unittest.skipIf(is_remote_only(), "Session creation different from local mode")
class SparkConnectSessionTests(ReusedConnectTestCase):
    def setUp(self) -> None:
        self.spark = (
            PySparkSession.builder.config(conf=self.conf())
            .appName(self.__class__.__name__)
            .remote("local[4]")
            .getOrCreate()
        )

    def tearDown(self):
        self.spark.stop()

    @timeout(10)
    def test_progress_handler(self):
        handler_called = []

        def handler(**kwargs):
            nonlocal handler_called
            handler_called.append(kwargs)

        self.spark.registerProgressHandler(handler)
        self.spark.sql("select 1").collect()
        self.assertGreaterEqual(len(handler_called), 1)

        handler_called = []
        self.spark.removeProgressHandler(handler)
        self.spark.sql("select 1").collect()
        self.assertEqual(len(handler_called), 0)

        self.spark.registerProgressHandler(handler)
        self.spark.clearProgressHandlers()
        self.spark.sql("select 1").collect()
        self.assertGreaterEqual(len(handler_called), 0)

    @timeout(10)
    def test_operation_id_in_execution_info_and_exception(self):
        df = self.spark.sql("select 1")
        df.collect()
        self.assertIsNotNone(df.executionInfo)
        operation_id = df.executionInfo.operation_id
        self.assertIsNotNone(operation_id)
        uuid.UUID(operation_id)

        with self.assertRaises(SparkConnectException) as error:
            self.spark.sql("select raise_error('expected')").collect()
        self.assertIsNotNone(error.exception.operation_id)
        uuid.UUID(error.exception.operation_id)

    def _check_no_active_session_error(self, e: PySparkException):
        self.check_error(exception=e, errorClass="NO_ACTIVE_SESSION", messageParameters=dict())

    @timeout(10)
    def test_stop_session(self):
        df = self.spark.sql("select 1 as a, 2 as b")
        catalog = self.spark.catalog
        self.spark.stop()

        # _execute_and_fetch
        with self.assertRaises(SparkConnectException) as e:
            self.spark.sql("select 1")
        self._check_no_active_session_error(e.exception)

        with self.assertRaises(SparkConnectException) as e:
            catalog.tableExists("table")
        self._check_no_active_session_error(e.exception)

        # _execute
        with self.assertRaises(SparkConnectException) as e:
            self.spark.udf.register("test_func", lambda x: x + 1)
        self._check_no_active_session_error(e.exception)

        # _analyze
        with self.assertRaises(SparkConnectException) as e:
            df._explain_string(extended=True)
        self._check_no_active_session_error(e.exception)

        # Config
        with self.assertRaises(SparkConnectException) as e:
            self.spark.conf.get("some.conf")
        self._check_no_active_session_error(e.exception)

    def test_with_block(self):
        with PySparkSession.builder.remote("local[*]").getOrCreate() as session:
            df = session.sql("select 1 as a, 2 as b")
            self.assertEqual(df.collect(), [(1, 2)])

        # Confirm that the session is stopped
        with self.assertRaises(SparkConnectException) as e:
            self.spark.sql("select 1")
        self._check_no_active_session_error(e.exception)

    def test_error_enrichment_message(self):
        with self.sql_conf(
            {
                "spark.sql.connect.enrichError.enabled": True,
                "spark.sql.connect.serverStacktrace.enabled": False,
                "spark.sql.pyspark.jvmStacktrace.enabled": False,
            }
        ):
            name = "test" * 10000
            with self.assertRaises(AnalysisException) as e:
                self.spark.sql("select " + name).collect()
            self.assertTrue(name in e.exception._message)
            self.assertFalse("JVM stacktrace" in e.exception._message)

    def test_error_enrichment_jvm_stacktrace(self):
        with self.sql_conf(
            {
                "spark.sql.connect.enrichError.enabled": True,
                "spark.sql.pyspark.jvmStacktrace.enabled": False,
                "spark.sql.legacy.timeParserPolicy": "EXCEPTION",
            }
        ):
            with self.sql_conf({"spark.sql.connect.serverStacktrace.enabled": False}):
                with self.assertRaises(SparkUpgradeException) as e:
                    self.spark.sql("""select from_json(
                            '{"d": "02-29"}', 'd date', map('dateFormat', 'MM-dd'))""").collect()
                self.assertFalse("JVM stacktrace" in e.exception._message)

            with self.sql_conf({"spark.sql.connect.serverStacktrace.enabled": True}):
                with self.assertRaises(SparkUpgradeException) as e:
                    self.spark.sql("""select from_json(
                            '{"d": "02-29"}', 'd date', map('dateFormat', 'MM-dd'))""").collect()
                self.assertTrue("JVM stacktrace" in str(e.exception))
                self.assertTrue("org.apache.spark.SparkUpgradeException" in str(e.exception))
                self.assertTrue(
                    "at org.apache.spark.sql.errors.ExecutionErrors"
                    ".failToParseDateTimeInNewParserError" in str(e.exception)
                )
                self.assertTrue("Caused by: java.time.DateTimeException:" in str(e.exception))

    def test_not_hitting_netty_header_limit(self):
        with self.sql_conf({"spark.sql.pyspark.jvmStacktrace.enabled": True}):
            with self.assertRaises(AnalysisException):
                self.spark.sql("select " + "test" * 1).collect()

    def test_error_stack_trace(self):
        with self.sql_conf({"spark.sql.connect.enrichError.enabled": False}):
            with self.sql_conf({"spark.sql.pyspark.jvmStacktrace.enabled": True}):
                with self.assertRaises(AnalysisException) as e:
                    self.spark.sql("select x").collect()
                self.assertTrue("JVM stacktrace" in str(e.exception))
                self.assertIsNotNone(e.exception.getStackTrace())
                self.assertTrue(
                    "at org.apache.spark.sql.catalyst.analysis.CheckAnalysis" in str(e.exception)
                )

            with self.sql_conf({"spark.sql.pyspark.jvmStacktrace.enabled": False}):
                with self.assertRaises(AnalysisException) as e:
                    self.spark.sql("select x").collect()
                self.assertFalse("JVM stacktrace" in str(e.exception))
                self.assertIsNone(e.exception.getStackTrace())
                self.assertFalse(
                    "at org.apache.spark.sql.catalyst.analysis.CheckAnalysis" in str(e.exception)
                )

        # Create a new session with a different stack trace size.
        self.spark.stop()
        spark = (
            PySparkSession.builder.config(conf=self.conf())
            .config("spark.connect.grpc.maxMetadataSize", 128)
            .remote("local[4]")
            .getOrCreate()
        )
        spark.conf.set("spark.sql.connect.enrichError.enabled", False)
        spark.conf.set("spark.sql.pyspark.jvmStacktrace.enabled", True)
        with self.assertRaises(AnalysisException) as e:
            spark.sql("select x").collect()
        self.assertTrue("JVM stacktrace" in str(e.exception))
        self.assertIsNotNone(e.exception.getStackTrace())
        self.assertFalse(
            "at org.apache.spark.sql.catalyst.analysis.CheckAnalysis" in str(e.exception)
        )
        spark.stop()

    def test_can_create_multiple_sessions_to_different_remotes(self):
        self.spark.stop()
        self.assertIsNotNone(self.spark._client)
        # Creates a new remote session.
        other = PySparkSession.builder.remote("sc://other.remote:114/").create()
        self.assertNotEqual(self.spark, other)

        # Gets currently active session.
        same = PySparkSession.builder.remote("sc://other.remote.host:114/").getOrCreate()
        self.assertEqual(other, same)
        same.release_session_on_close = False  # avoid sending release to dummy connection
        same.stop()

        # Make sure the environment is clean.
        self.spark.stop()
        with self.assertRaises(RuntimeError) as e:
            PySparkSession.builder.create()
            self.assertIn("Create a new SparkSession is only supported with SparkConnect.", str(e))

    def test_get_message_parameters_without_enriched_error(self):
        with self.sql_conf({"spark.sql.connect.enrichError.enabled": False}):
            exception = None
            try:
                self.spark.sql("""SELECT a""")
            except AnalysisException as e:
                exception = e

            self.assertIsNotNone(exception)
            self.assertEqual(exception.getMessageParameters(), {"objectName": "`a`"})

    @timeout(10)
    def test_custom_channel_builder(self):
        # Access self.spark's DefaultChannelBuilder to reuse same endpoint
        endpoint = self.spark._client._builder.endpoint

        class CustomChannelBuilder(ChannelBuilder):
            def toChannel(self):
                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(endpoint, creds)

        session = RemoteSparkSession.builder.channelBuilder(CustomChannelBuilder()).create()
        session.sql("select 1 + 1")

    def test_reset_when_server_and_client_sessionids_mismatch(self):
        session = RemoteSparkSession.builder.remote("sc://localhost").getOrCreate()
        # run a simple query so the session id is synchronized.
        session.range(3).collect()

        # trigger a mismatch between client session id and server session id.
        session._client._session_id = str(uuid.uuid4())
        with self.assertRaises(SparkConnectException):
            session.range(3).collect()

        # assert that getOrCreate() generates a new session
        session = RemoteSparkSession.builder.remote("sc://localhost").getOrCreate()
        session.range(3).collect()

    def test_reset_when_server_session_id_mismatch(self):
        session = RemoteSparkSession.builder.remote("sc://localhost").getOrCreate()
        # run a simple query so the session id is synchronized.
        session.range(3).collect()

        # trigger a mismatch
        session._client._server_session_id = str(uuid.uuid4())
        with self.assertRaises(SparkConnectException):
            session.range(3).collect()

        # assert that getOrCreate() generates a new session
        session = RemoteSparkSession.builder.remote("sc://localhost").getOrCreate()
        session.range(3).collect()

    def test_stop_invalid_session(self):  # SPARK-47986
        session = RemoteSparkSession.builder.remote("sc://localhost").getOrCreate()
        # run a simple query so the session id is synchronized.
        session.range(3).collect()

        # change the server side session id to simulate that the server has terminated this session.
        session._client._server_session_id = str(uuid.uuid4())

        # Should not throw any error
        session.stop()

    def test_api_mode(self):
        session = (
            PySparkSession.builder.config("spark.api.mode", "connect")
            .master("sc://localhost")
            .getOrCreate()
        )
        self.assertEqual(session.range(1).first()[0], 0)
        self.assertIsInstance(session, RemoteSparkSession)

    def test_authentication(self):
        # All servers start with a default token of "deadbeef", so supply in invalid one
        session = RemoteSparkSession.builder.remote("sc://localhost/;token=invalid").create()

        with self.assertRaises(SparkConnectGrpcException) as e:
            session.range(3).collect()

        self.assertTrue("Invalid authentication token" in str(e.exception))


@unittest.skipIf(not should_test_connect, connect_requirement_message)
class SparkConnectSessionBuilderTests(unittest.TestCase):
    def setUp(self):
        # Reset class-level session state so tests are order-independent.
        RemoteSparkSession._default_session = None
        RemoteSparkSession._active_session.session = None

    def test_fails_to_create_session_without_remote_and_channel_builder(self):
        with self.assertRaises(ValueError):
            RemoteSparkSession.builder.getOrCreate()

    def test_fails_to_create_when_both_remote_and_channel_builder_are_specified(self):
        with self.assertRaises(ValueError):
            (
                RemoteSparkSession.builder.channelBuilder(CustomChannelBuilder("sc://localhost"))
                .remote("sc://localhost")
                .getOrCreate()
            )

    def test_creates_session_with_channel_builder(self):
        test_session = RemoteSparkSession.builder.channelBuilder(
            CustomChannelBuilder("sc://other")
        ).getOrCreate()
        host = test_session.client.host
        # Skip release_session() since "sc://other" is a fake remote.
        test_session.release_session_on_close = False
        test_session.stop()

        self.assertEqual("other", host)

    def test_creates_session_with_remote(self):
        test_session = RemoteSparkSession.builder.remote("sc://other").getOrCreate()
        host = test_session.client.host
        test_session.release_session_on_close = False
        test_session.stop()

        self.assertEqual("other", host)

    def test_session_stop(self):
        session = RemoteSparkSession.builder.remote("sc://other").getOrCreate()
        session.release_session_on_close = False

        self.assertFalse(session.is_stopped)
        session.stop()
        self.assertTrue(session.is_stopped)

    def test_session_create_sets_active_session(self):
        session = RemoteSparkSession.builder.remote("sc://abc").create()
        session2 = RemoteSparkSession.builder.remote("sc://other").getOrCreate()

        self.assertIs(session, session2)
        session.release_session_on_close = False
        session.stop()

    def test_active_session_expires_when_client_closes(self):
        s1 = RemoteSparkSession.builder.remote("sc://other").getOrCreate()
        s2 = RemoteSparkSession.getActiveSession()

        self.assertIs(s1, s2)

        # We don't call close() to avoid executing ExecutePlanResponseReattachableIterator
        s1._client._closed = True

        self.assertIsNone(RemoteSparkSession.getActiveSession())
        s3 = RemoteSparkSession.builder.remote("sc://other").getOrCreate()

        self.assertIsNot(s1, s3)

    def test_default_session_expires_when_client_closes(self):
        s1 = RemoteSparkSession.builder.remote("sc://other").getOrCreate()
        s2 = RemoteSparkSession._get_default_session()

        self.assertIs(s1, s2)

        # We don't call close() to avoid executing ExecutePlanResponseReattachableIterator
        s1._client._closed = True

        self.assertIsNone(RemoteSparkSession._get_default_session())
        s3 = RemoteSparkSession.builder.remote("sc://other").getOrCreate()

        self.assertIsNot(s1, s3)


@unittest.skipIf(not should_test_connect, connect_requirement_message)
class SparkConnectSessionWithOptionsTest(unittest.TestCase):
    def setUp(self) -> None:
        self.spark = (
            PySparkSession.builder.config("string", "foo")
            .config("integer", 1)
            .config("boolean", False)
            .appName(self.__class__.__name__)
            .remote(os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[4]"))
            .getOrCreate()
        )

    def tearDown(self):
        self.spark.stop()

    def test_config(self):
        # Config
        self.assertEqual(self.spark.conf.get("string"), "foo")
        self.assertEqual(self.spark.conf.get("boolean"), "false")
        self.assertEqual(self.spark.conf.get("integer"), "1")


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

    main()