You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by hv...@apache.org on 2025/10/13 15:02:23 UTC
(spark) branch master updated: [SPARK-53455][CONNECT] Add `CloneSession` RPC
This is an automated email from the ASF dual-hosted git repository.
hvanhovell pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 264ca4dc320e [SPARK-53455][CONNECT] Add `CloneSession` RPC
264ca4dc320e is described below
commit 264ca4dc320e516bd88679ed2404f1c15d2f24eb
Author: vicennial <ve...@databricks.com>
AuthorDate: Mon Oct 13 11:01:58 2025 -0400
[SPARK-53455][CONNECT] Add `CloneSession` RPC
### What changes were proposed in this pull request?
Adds a new experimental/developer RPC `CloneSession` to the `SparkConnectService`.
✅ CLONED (from SparkSession)
- SessionState - SQL configs, temp views, UDFs, catalog metadata
- ArtifactManager - JARs, files, classes added to session
- ManagedJobTags - Job group tags for tracking
- SharedState (reference) - Metastore, global temp views
- SparkContext (reference) - Core Spark engine
❌ NOT CLONED (SessionHolder - Spark Connect layer)
- planCache - (Partially analyzed) Logical plans for query optimization
- operationIds - Currently executing operations
- errorIdToError - Recent errors for debugging
- eventManager - Session lifecycle events
- dataFrameCache - DataFrames for foreachBatch callbacks
- mlCache - ML models and pipelines
- listenerCache - Streaming query listeners
- pipelineExecutions - Active pipeline contexts
- dataflowGraphRegistry - Registered dataflow graphs
- streamingForeachBatchRunnerCleanerCache - Python streaming workers
- pythonAccumulator - Python metrics collection
- Session timings - Start time, last access, custom timeout
The clone preserves all SQL/catalog state but creates a fresh runtime environment. An analogy is cloning a database schema/config but not the active connections, caches, or running jobs.
### Why are the changes needed?
Spark Connect introduced the concept of resource isolation (via `ArtifactManager`, which has been ported to classic Spark) and thus, jars/pyfiles/artifacts added to each session are isolated from other sessions.
A slight rough edge is that if a user wishes to fork the state of a session but maintain independence, the only possible way is to create a new session and reupload/reinit all base jars/artifacts/pyfiles, etc.
Support for cloning through the API helps address the rough edge while maintaining all the benefits of session resource isolation.
### Does this PR introduce _any_ user-facing change?
Yes
```python
spark = SparkSession.builder.remote("sc://localhost").getOrCreate()
spark.conf.set("my.custom.config", "value")
spark.addArtifact("/path/to/my.jar")
spark.sql("CREATE TEMP VIEW my_view AS SELECT 1 AS id")
# Clone the session
cloned_spark = spark.cloneSession()
# The cloned session has all the same state
assert cloned_spark.conf.get("my.custom.config") == "value"
assert cloned_spark.sql("SELECT * FROM my_view").collect() == [Row(id=1)]
# But operations are isolated between sessions
cloned_spark.sql("DROP VIEW my_view") # Only affects cloned session
spark.sql("SELECT * FROM my_view").collect() # Original still works
```
### How was this patch tested?
New individual unit tests along with new test suites.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with assistance from Claude Code.
Closes #52200 from vicennial/cloneAPI.
Authored-by: vicennial <ve...@databricks.com>
Signed-off-by: Herman van Hovell <he...@databricks.com>
---
.../src/main/resources/error/error-conditions.json | 23 ++
dev/sparktestsupport/modules.py | 1 +
python/pyspark/sql/connect/client/core.py | 66 +++++
python/pyspark/sql/connect/proto/base_pb2.py | 10 +-
python/pyspark/sql/connect/proto/base_pb2.pyi | 149 +++++++++++
python/pyspark/sql/connect/proto/base_pb2_grpc.py | 55 ++++
python/pyspark/sql/connect/session.py | 34 +++
.../tests/connect/test_connect_clone_session.py | 153 +++++++++++
.../sql/tests/test_connect_compatibility.py | 2 +-
.../spark/sql/connect/SparkSessionCloneSuite.scala | 64 +++++
.../src/main/protobuf/spark/connect/base.proto | 58 +++++
.../apache/spark/sql/connect/SparkSession.scala | 33 +++
.../client/CustomSparkConnectBlockingStub.scala | 13 +
.../sql/connect/client/SparkConnectClient.scala | 51 ++++
.../service/SparkConnectCloneSessionHandler.scala | 62 +++++
.../sql/connect/service/SparkConnectService.scala | 14 ++
.../service/SparkConnectSessionManager.scala | 67 +++++
.../service/SparkConnectCloneSessionSuite.scala | 279 +++++++++++++++++++++
18 files changed, 1130 insertions(+), 4 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json
index 8f6687587f78..afe56f6db2fd 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -2507,6 +2507,29 @@
],
"sqlState" : "22P03"
},
+ "INVALID_CLONE_SESSION_REQUEST" : {
+ "message" : [
+ "Invalid session clone request."
+ ],
+ "subClass" : {
+ "TARGET_SESSION_ID_ALREADY_CLOSED" : {
+ "message" : [
+ "Cannot clone session to target session ID <targetSessionId> because a session with this ID was previously closed."
+ ]
+ },
+ "TARGET_SESSION_ID_ALREADY_EXISTS" : {
+ "message" : [
+ "Cannot clone session to target session ID <targetSessionId> because a session with this ID already exists."
+ ]
+ },
+ "TARGET_SESSION_ID_FORMAT" : {
+ "message" : [
+ "Target session ID <targetSessionId> for clone operation must be an UUID string of the format '00112233-4455-6677-8899-aabbccddeeff'."
+ ]
+ }
+ },
+ "sqlState" : "42K04"
+ },
"INVALID_COLUMN_NAME_AS_PATH" : {
"message" : [
"The datasource <datasource> cannot save the column <columnName> because its name contains some characters that are not allowed in file paths. Please, use an alias to rename it."
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index 8aab600071dc..945a2ac9189b 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -1095,6 +1095,7 @@ pyspark_connect = Module(
"pyspark.sql.tests.connect.test_connect_basic",
"pyspark.sql.tests.connect.test_connect_dataframe_property",
"pyspark.sql.tests.connect.test_connect_channel",
+ "pyspark.sql.tests.connect.test_connect_clone_session",
"pyspark.sql.tests.connect.test_connect_error",
"pyspark.sql.tests.connect.test_connect_function",
"pyspark.sql.tests.connect.test_connect_collection",
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index 741d612f53f4..d0d191dbd7fd 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -2105,3 +2105,69 @@ class SparkConnectClient(object):
ml_command_result = properties["ml_command_result"]
return ml_command_result.param.long
+
+ 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
diff --git a/python/pyspark/sql/connect/proto/base_pb2.py b/python/pyspark/sql/connect/proto/base_pb2.py
index 269f5eefbd26..069e0830fcea 100644
--- a/python/pyspark/sql/connect/proto/base_pb2.py
+++ b/python/pyspark/sql/connect/proto/base_pb2.py
@@ -45,7 +45,7 @@ from pyspark.sql.connect.proto import pipelines_pb2 as spark_dot_connect_dot_pip
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
- b'\n\x18spark/connect/base.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1cspark/connect/commands.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\x1a\x19spark/connect/types.proto\x1a\x16spark/connect/ml.proto\x1a\x1dspark/connect/pipelines.proto"t\n\x04Plan\x12-\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x04root\x12\x32\n\x07\x63ommand\x18\x02 \x01(\x0b\x32\x16.spark.connect.Com [...]
+ b'\n\x18spark/connect/base.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1cspark/connect/commands.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\x1a\x19spark/connect/types.proto\x1a\x16spark/connect/ml.proto\x1a\x1dspark/connect/pipelines.proto"t\n\x04Plan\x12-\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x04root\x12\x32\n\x07\x63ommand\x18\x02 \x01(\x0b\x32\x16.spark.connect.Com [...]
)
_globals = globals()
@@ -260,6 +260,10 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_end = 17518
_globals["_CHECKPOINTCOMMANDRESULT"]._serialized_start = 17539
_globals["_CHECKPOINTCOMMANDRESULT"]._serialized_end = 17629
- _globals["_SPARKCONNECTSERVICE"]._serialized_start = 17632
- _globals["_SPARKCONNECTSERVICE"]._serialized_end = 18578
+ _globals["_CLONESESSIONREQUEST"]._serialized_start = 17632
+ _globals["_CLONESESSIONREQUEST"]._serialized_end = 17994
+ _globals["_CLONESESSIONRESPONSE"]._serialized_start = 17997
+ _globals["_CLONESESSIONRESPONSE"]._serialized_end = 18201
+ _globals["_SPARKCONNECTSERVICE"]._serialized_start = 18204
+ _globals["_SPARKCONNECTSERVICE"]._serialized_end = 19241
# @@protoc_insertion_point(module_scope)
diff --git a/python/pyspark/sql/connect/proto/base_pb2.pyi b/python/pyspark/sql/connect/proto/base_pb2.pyi
index 3361649d9323..dc3099ecdffc 100644
--- a/python/pyspark/sql/connect/proto/base_pb2.pyi
+++ b/python/pyspark/sql/connect/proto/base_pb2.pyi
@@ -4081,3 +4081,152 @@ class CheckpointCommandResult(google.protobuf.message.Message):
) -> None: ...
global___CheckpointCommandResult = CheckpointCommandResult
+
+class CloneSessionRequest(google.protobuf.message.Message):
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+ SESSION_ID_FIELD_NUMBER: builtins.int
+ CLIENT_OBSERVED_SERVER_SIDE_SESSION_ID_FIELD_NUMBER: builtins.int
+ USER_CONTEXT_FIELD_NUMBER: builtins.int
+ CLIENT_TYPE_FIELD_NUMBER: builtins.int
+ NEW_SESSION_ID_FIELD_NUMBER: builtins.int
+ session_id: builtins.str
+ """(Required)
+
+ The session_id specifies a spark session for a user id (which is specified
+ by user_context.user_id). The session_id is set by the client to be able to
+ collate streaming responses from different queries within the dedicated session.
+ The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff`
+ """
+ client_observed_server_side_session_id: builtins.str
+ """(Optional)
+
+ Server-side generated idempotency key from the previous responses (if any). Server
+ can use this to validate that the server side session has not changed.
+ """
+ @property
+ def user_context(self) -> global___UserContext:
+ """(Required) User context
+
+ user_context.user_id and session_id both identify a unique remote spark session on the
+ server side.
+ """
+ client_type: builtins.str
+ """Provides optional information about the client sending the request. This field
+ can be used for language or version specific information and is only intended for
+ logging purposes and will not be interpreted by the server.
+ """
+ new_session_id: builtins.str
+ """(Optional)
+ The session_id for the new cloned session. If not provided, a new UUID will be generated.
+ The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff`
+ """
+ def __init__(
+ self,
+ *,
+ session_id: builtins.str = ...,
+ client_observed_server_side_session_id: builtins.str | None = ...,
+ user_context: global___UserContext | None = ...,
+ client_type: builtins.str | None = ...,
+ new_session_id: builtins.str | None = ...,
+ ) -> None: ...
+ def HasField(
+ self,
+ field_name: typing_extensions.Literal[
+ "_client_observed_server_side_session_id",
+ b"_client_observed_server_side_session_id",
+ "_client_type",
+ b"_client_type",
+ "_new_session_id",
+ b"_new_session_id",
+ "client_observed_server_side_session_id",
+ b"client_observed_server_side_session_id",
+ "client_type",
+ b"client_type",
+ "new_session_id",
+ b"new_session_id",
+ "user_context",
+ b"user_context",
+ ],
+ ) -> builtins.bool: ...
+ def ClearField(
+ self,
+ field_name: typing_extensions.Literal[
+ "_client_observed_server_side_session_id",
+ b"_client_observed_server_side_session_id",
+ "_client_type",
+ b"_client_type",
+ "_new_session_id",
+ b"_new_session_id",
+ "client_observed_server_side_session_id",
+ b"client_observed_server_side_session_id",
+ "client_type",
+ b"client_type",
+ "new_session_id",
+ b"new_session_id",
+ "session_id",
+ b"session_id",
+ "user_context",
+ b"user_context",
+ ],
+ ) -> None: ...
+ @typing.overload
+ def WhichOneof(
+ self,
+ oneof_group: typing_extensions.Literal[
+ "_client_observed_server_side_session_id", b"_client_observed_server_side_session_id"
+ ],
+ ) -> typing_extensions.Literal["client_observed_server_side_session_id"] | None: ...
+ @typing.overload
+ def WhichOneof(
+ self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"]
+ ) -> typing_extensions.Literal["client_type"] | None: ...
+ @typing.overload
+ def WhichOneof(
+ self, oneof_group: typing_extensions.Literal["_new_session_id", b"_new_session_id"]
+ ) -> typing_extensions.Literal["new_session_id"] | None: ...
+
+global___CloneSessionRequest = CloneSessionRequest
+
+class CloneSessionResponse(google.protobuf.message.Message):
+ """Next ID: 5"""
+
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+ SESSION_ID_FIELD_NUMBER: builtins.int
+ SERVER_SIDE_SESSION_ID_FIELD_NUMBER: builtins.int
+ NEW_SESSION_ID_FIELD_NUMBER: builtins.int
+ NEW_SERVER_SIDE_SESSION_ID_FIELD_NUMBER: builtins.int
+ session_id: builtins.str
+ """Session id of the original session that was cloned."""
+ server_side_session_id: builtins.str
+ """Server-side generated idempotency key that the client can use to assert that the server side
+ session (parent of the cloned session) has not changed.
+ """
+ new_session_id: builtins.str
+ """Session id of the new cloned session."""
+ new_server_side_session_id: builtins.str
+ """Server-side session ID of the new cloned session."""
+ def __init__(
+ self,
+ *,
+ session_id: builtins.str = ...,
+ server_side_session_id: builtins.str = ...,
+ new_session_id: builtins.str = ...,
+ new_server_side_session_id: builtins.str = ...,
+ ) -> None: ...
+ def ClearField(
+ self,
+ field_name: typing_extensions.Literal[
+ "new_server_side_session_id",
+ b"new_server_side_session_id",
+ "new_session_id",
+ b"new_session_id",
+ "server_side_session_id",
+ b"server_side_session_id",
+ "session_id",
+ b"session_id",
+ ],
+ ) -> None: ...
+
+global___CloneSessionResponse = CloneSessionResponse
diff --git a/python/pyspark/sql/connect/proto/base_pb2_grpc.py b/python/pyspark/sql/connect/proto/base_pb2_grpc.py
index 7501aaf0a3a2..16dd20b563f3 100644
--- a/python/pyspark/sql/connect/proto/base_pb2_grpc.py
+++ b/python/pyspark/sql/connect/proto/base_pb2_grpc.py
@@ -90,6 +90,12 @@ class SparkConnectServiceStub(object):
response_deserializer=spark_dot_connect_dot_base__pb2.FetchErrorDetailsResponse.FromString,
_registered_method=True,
)
+ self.CloneSession = channel.unary_unary(
+ "/spark.connect.SparkConnectService/CloneSession",
+ request_serializer=spark_dot_connect_dot_base__pb2.CloneSessionRequest.SerializeToString,
+ response_deserializer=spark_dot_connect_dot_base__pb2.CloneSessionResponse.FromString,
+ _registered_method=True,
+ )
class SparkConnectServiceServicer(object):
@@ -172,6 +178,20 @@ class SparkConnectServiceServicer(object):
context.set_details("Method not implemented!")
raise NotImplementedError("Method not implemented!")
+ def CloneSession(self, request, context):
+ """Create a clone of a Spark Connect 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 cloned session
+ is isolated from the source session - any subsequent changes to either session's
+ server-side state will not be reflected in the other.
+
+ The request can optionally specify a custom session ID for the cloned session (must be
+ a valid UUID). If not provided, a new UUID will be generated automatically.
+ """
+ context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+ context.set_details("Method not implemented!")
+ raise NotImplementedError("Method not implemented!")
+
def add_SparkConnectServiceServicer_to_server(servicer, server):
rpc_method_handlers = {
@@ -225,6 +245,11 @@ def add_SparkConnectServiceServicer_to_server(servicer, server):
request_deserializer=spark_dot_connect_dot_base__pb2.FetchErrorDetailsRequest.FromString,
response_serializer=spark_dot_connect_dot_base__pb2.FetchErrorDetailsResponse.SerializeToString,
),
+ "CloneSession": grpc.unary_unary_rpc_method_handler(
+ servicer.CloneSession,
+ request_deserializer=spark_dot_connect_dot_base__pb2.CloneSessionRequest.FromString,
+ response_serializer=spark_dot_connect_dot_base__pb2.CloneSessionResponse.SerializeToString,
+ ),
}
generic_handler = grpc.method_handlers_generic_handler(
"spark.connect.SparkConnectService", rpc_method_handlers
@@ -536,3 +561,33 @@ class SparkConnectService(object):
metadata,
_registered_method=True,
)
+
+ @staticmethod
+ def CloneSession(
+ request,
+ target,
+ options=(),
+ channel_credentials=None,
+ call_credentials=None,
+ insecure=False,
+ compression=None,
+ wait_for_ready=None,
+ timeout=None,
+ metadata=None,
+ ):
+ return grpc.experimental.unary_unary(
+ request,
+ target,
+ "/spark.connect.SparkConnectService/CloneSession",
+ spark_dot_connect_dot_base__pb2.CloneSessionRequest.SerializeToString,
+ spark_dot_connect_dot_base__pb2.CloneSessionResponse.FromString,
+ options,
+ channel_credentials,
+ insecure,
+ call_credentials,
+ compression,
+ wait_for_ready,
+ timeout,
+ metadata,
+ _registered_method=True,
+ )
diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py
index 6ccffc718d06..f759137fac1d 100644
--- a/python/pyspark/sql/connect/session.py
+++ b/python/pyspark/sql/connect/session.py
@@ -1175,6 +1175,40 @@ class SparkSession:
assert dt is not None
return dt
+ def cloneSession(self, new_session_id: Optional[str] = None) -> "SparkSession":
+ """
+ Create a clone of this Spark Connect 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 cloned
+ session is isolated from this 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
+ -------
+ SparkSession
+ A new SparkSession 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.
+ """
+ cloned_client = self._client.clone(new_session_id)
+ # Create a new SparkSession with the cloned client directly
+ new_session = object.__new__(SparkSession)
+ new_session._client = cloned_client
+ new_session._session_id = cloned_client._session_id
+ return new_session
+
SparkSession.__doc__ = PySparkSession.__doc__
diff --git a/python/pyspark/sql/tests/connect/test_connect_clone_session.py b/python/pyspark/sql/tests/connect/test_connect_clone_session.py
new file mode 100644
index 000000000000..1cef628a158c
--- /dev/null
+++ b/python/pyspark/sql/tests/connect/test_connect_clone_session.py
@@ -0,0 +1,153 @@
+#
+# 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 uuid
+import unittest
+
+from pyspark.sql.tests.connect.test_connect_basic import SparkConnectSQLTestCase
+
+
+class SparkConnectCloneSessionTest(SparkConnectSQLTestCase):
+ def test_clone_session_basic(self):
+ """Test basic session cloning functionality."""
+ # Set a configuration in the original session
+ self.connect.sql("SET spark.test.original = 'value1'")
+
+ # Verify the original session has the value
+ original_value = self.connect.sql("SET spark.test.original").collect()[0][1]
+ self.assertEqual(original_value, "'value1'")
+
+ # Clone the session
+ cloned_session = self.connect.cloneSession()
+
+ # Verify the configuration was copied
+ # (if cloning doesn't preserve dynamic configs, use a different approach)
+ cloned_value = cloned_session.sql("SET spark.test.original").collect()[0][1]
+ self.assertEqual(cloned_value, "'value1'")
+ # Verify that sessions are independent by setting different values
+ cloned_session.sql("SET spark.test.original = 'modified_clone'")
+ self.connect.sql("SET spark.test.original = 'modified_original'")
+
+ # Verify independence
+ original_final = self.connect.sql("SET spark.test.original").collect()[0][1]
+ cloned_final = cloned_session.sql("SET spark.test.original").collect()[0][1]
+
+ self.assertEqual(original_final, "'modified_original'")
+ self.assertEqual(cloned_final, "'modified_clone'")
+
+ def test_clone_session_with_custom_id(self):
+ """Test cloning session with a custom session ID."""
+ custom_session_id = str(uuid.uuid4())
+ cloned_session = self.connect.cloneSession(custom_session_id)
+
+ self.assertEqual(cloned_session.session_id, custom_session_id)
+ self.assertNotEqual(self.connect.session_id, cloned_session.session_id)
+
+ def test_clone_session_preserves_temp_views(self):
+ """Test that temporary views are preserved in cloned sessions."""
+ # Create a temporary view
+ df = self.connect.createDataFrame([(1, "a"), (2, "b")], ["id", "value"])
+ df.createOrReplaceTempView("test_table")
+
+ # Verify original session can access the temp view
+ original_count = self.connect.sql("SELECT COUNT(*) FROM test_table").collect()[0][0]
+ self.assertEqual(original_count, 2)
+
+ # Clone the session
+ cloned_session = self.connect.cloneSession()
+
+ # Create temp view in cloned session for testing
+ df_clone = cloned_session.createDataFrame([(3, "c"), (4, "d")], ["id", "value"])
+ df_clone.createOrReplaceTempView("test_table")
+
+ # Verify both sessions have independent temp views
+ original_count_final = self.connect.sql("SELECT COUNT(*) FROM test_table").collect()[0][0]
+ cloned_count_final = cloned_session.sql("SELECT COUNT(*) FROM test_table").collect()[0][0]
+
+ self.assertEqual(original_count_final, 2) # Original data
+ self.assertEqual(cloned_count_final, 2) # New data
+
+ def test_temp_views_independence_after_cloning(self):
+ """Test that temp views are cloned and then can be modified independently."""
+ # Create initial temp view before cloning
+ df = self.connect.createDataFrame([(1, "original")], ["id", "type"])
+ df.createOrReplaceTempView("shared_view")
+
+ # Verify original session can access the temp view
+ original_count_before = self.connect.sql("SELECT COUNT(*) FROM shared_view").collect()[0][0]
+ self.assertEqual(original_count_before, 1)
+
+ # Clone the session - temp views should be preserved
+ cloned_session = self.connect.cloneSession()
+
+ # Verify cloned session can access the same temp view (cloned)
+ cloned_count = cloned_session.sql("SELECT COUNT(*) FROM shared_view").collect()[0][0]
+ self.assertEqual(cloned_count, 1)
+
+ # Now modify the temp view in each session independently
+ # Replace the view in the original session
+ original_df = self.connect.createDataFrame(
+ [(2, "modified_original"), (3, "another")], ["id", "type"]
+ )
+ original_df.createOrReplaceTempView("shared_view")
+
+ # Replace the view in the cloned session
+ cloned_session.sql(
+ "CREATE OR REPLACE TEMPORARY VIEW shared_view AS SELECT 4 as id, 'cloned_data' as type"
+ )
+
+ # Verify they now have different content (independence after cloning)
+ original_count_after = self.connect.sql("SELECT COUNT(*) FROM shared_view").collect()[0][0]
+ cloned_count_after = cloned_session.sql("SELECT COUNT(*) FROM shared_view").collect()[0][0]
+
+ self.assertEqual(original_count_after, 2) # Original session: 2 rows
+ self.assertEqual(cloned_count_after, 1) # Cloned session: 1 row
+
+ def test_invalid_session_id_format(self):
+ """Test that invalid session ID format raises an exception."""
+ with self.assertRaises(Exception) as context:
+ self.connect.cloneSession("not-a-valid-uuid")
+
+ # Verify it contains our clone-specific error message
+ self.assertIn(
+ "INVALID_CLONE_SESSION_REQUEST.TARGET_SESSION_ID_FORMAT", str(context.exception)
+ )
+ self.assertIn("not-a-valid-uuid", str(context.exception))
+
+ def test_clone_session_auto_generated_id(self):
+ """Test that cloneSession() without arguments generates a valid UUID."""
+ cloned_session = self.connect.cloneSession()
+
+ # Verify different session IDs
+ self.assertNotEqual(self.connect.session_id, cloned_session.session_id)
+
+ # Verify the new session ID is a valid UUID
+ try:
+ uuid.UUID(cloned_session.session_id)
+ except ValueError:
+ self.fail("Generated session ID is not a valid UUID")
+
+
+if __name__ == "__main__":
+ from pyspark.sql.tests.connect.test_connect_clone_session import * # noqa: F401
+
+ try:
+ import xmlrunner
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/python/pyspark/sql/tests/test_connect_compatibility.py b/python/pyspark/sql/tests/test_connect_compatibility.py
index b2e0cc6229c4..4ca1c04bf2d0 100644
--- a/python/pyspark/sql/tests/test_connect_compatibility.py
+++ b/python/pyspark/sql/tests/test_connect_compatibility.py
@@ -268,7 +268,7 @@ class ConnectCompatibilityTestsMixin:
"registerProgressHandler",
"removeProgressHandler",
}
- expected_missing_classic_methods = set()
+ expected_missing_classic_methods = {"cloneSession"}
self.check_compatibility(
ClassicSparkSession,
ConnectSparkSession,
diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/SparkSessionCloneSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/SparkSessionCloneSuite.scala
new file mode 100644
index 000000000000..97ad21b2675f
--- /dev/null
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/SparkSessionCloneSuite.scala
@@ -0,0 +1,64 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.connect
+
+import java.util.UUID
+
+import org.apache.spark.sql.connect.test.{ConnectFunSuite, RemoteSparkSession}
+
+class SparkSessionCloneSuite extends ConnectFunSuite with RemoteSparkSession {
+
+ test("cloneSession() creates independent session with auto-generated ID") {
+ spark.sql("SET spark.test.original = 'value1'")
+
+ val clonedSession = spark.cloneSession()
+
+ // Sessions should have different IDs
+ assert(spark.sessionId !== clonedSession.sessionId)
+
+ // Modify original session setting
+ spark.sql("SET spark.test.original = 'modified_original'")
+
+ // Verify original session has the modified value
+ val originalFinal = spark.sql("SET spark.test.original").collect().head.getString(1)
+ assert(originalFinal === "'modified_original'")
+
+ // Verify cloned session retains the original value
+ val clonedValue = clonedSession.sql("SET spark.test.original").collect().head.getString(1)
+ assert(clonedValue === "'value1'")
+ }
+
+ test("cloneSession(sessionId) creates session with specified ID") {
+ val customSessionId = UUID.randomUUID().toString
+
+ val clonedSession = spark.cloneSession(customSessionId)
+
+ // Verify the cloned session has the specified ID
+ assert(clonedSession.sessionId === customSessionId)
+ assert(spark.sessionId !== clonedSession.sessionId)
+ }
+
+ test("invalid session ID format throws exception") {
+ val ex = intercept[org.apache.spark.SparkException] {
+ spark.cloneSession("not-a-valid-uuid")
+ }
+ // Verify it contains our clone-specific error message
+ assert(ex.getMessage.contains("INVALID_CLONE_SESSION_REQUEST.TARGET_SESSION_ID_FORMAT"))
+ assert(ex.getMessage.contains("not-a-valid-uuid"))
+ }
+}
diff --git a/sql/connect/common/src/main/protobuf/spark/connect/base.proto b/sql/connect/common/src/main/protobuf/spark/connect/base.proto
index f34b169f821d..6e1029bf0a6a 100644
--- a/sql/connect/common/src/main/protobuf/spark/connect/base.proto
+++ b/sql/connect/common/src/main/protobuf/spark/connect/base.proto
@@ -1152,6 +1152,54 @@ message CheckpointCommandResult {
CachedRemoteRelation relation = 1;
}
+message CloneSessionRequest {
+ // (Required)
+ //
+ // The session_id specifies a spark session for a user id (which is specified
+ // by user_context.user_id). The session_id is set by the client to be able to
+ // collate streaming responses from different queries within the dedicated session.
+ // The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff`
+ string session_id = 1;
+
+ // (Optional)
+ //
+ // Server-side generated idempotency key from the previous responses (if any). Server
+ // can use this to validate that the server side session has not changed.
+ optional string client_observed_server_side_session_id = 5;
+
+ // (Required) User context
+ //
+ // user_context.user_id and session_id both identify a unique remote spark session on the
+ // server side.
+ UserContext user_context = 2;
+
+ // Provides optional information about the client sending the request. This field
+ // can be used for language or version specific information and is only intended for
+ // logging purposes and will not be interpreted by the server.
+ optional string client_type = 3;
+
+ // (Optional)
+ // The session_id for the new cloned session. If not provided, a new UUID will be generated.
+ // The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff`
+ optional string new_session_id = 4;
+}
+
+// Next ID: 5
+message CloneSessionResponse {
+ // Session id of the original session that was cloned.
+ string session_id = 1;
+
+ // Server-side generated idempotency key that the client can use to assert that the server side
+ // session (parent of the cloned session) has not changed.
+ string server_side_session_id = 2;
+
+ // Session id of the new cloned session.
+ string new_session_id = 3;
+
+ // Server-side session ID of the new cloned session.
+ string new_server_side_session_id = 4;
+}
+
// Main interface for the SparkConnect service.
service SparkConnectService {
@@ -1196,4 +1244,14 @@ service SparkConnectService {
// FetchErrorDetails retrieves the matched exception with details based on a provided error id.
rpc FetchErrorDetails(FetchErrorDetailsRequest) returns (FetchErrorDetailsResponse) {}
+
+ // Create a clone of a Spark Connect 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 cloned session
+ // is isolated from the source session - any subsequent changes to either session's
+ // server-side state will not be reflected in the other.
+ //
+ // The request can optionally specify a custom session ID for the cloned session (must be
+ // a valid UUID). If not provided, a new UUID will be generated automatically.
+ rpc CloneSession(CloneSessionRequest) returns (CloneSessionResponse) {}
}
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/SparkSession.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/SparkSession.scala
index 646db83981fe..f7869a8b4dd8 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/SparkSession.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/SparkSession.scala
@@ -684,6 +684,39 @@ class SparkSession private[sql] (
}
}
+ /**
+ * Create a clone of this Spark Connect 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 cloned session is
+ * isolated from this session - any subsequent changes to either session's server-side state
+ * will not be reflected in the other.
+ *
+ * @note
+ * This creates a new server-side session with a new session ID while preserving the current
+ * session's configuration and state.
+ */
+ @DeveloperApi
+ def cloneSession(): SparkSession = {
+ SparkSession.builder().client(client.cloneSession()).create()
+ }
+
+ /**
+ * Create a clone of this Spark Connect session on the server side with a custom session ID. 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 with the
+ * specified session ID. The returned cloned session is isolated from this session.
+ *
+ * @param sessionId
+ * The custom session ID to use for the cloned session (must be a valid UUID)
+ * @note
+ * This creates a new server-side session with the specified session ID while preserving the
+ * current session's configuration and state.
+ */
+ @DeveloperApi
+ def cloneSession(sessionId: String): SparkSession = {
+ SparkSession.builder().client(client.clone(sessionId)).create()
+ }
+
override private[sql] def isUsable: Boolean = client.isSessionValid
}
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala
index d7867229248b..913f068fcf34 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala
@@ -131,4 +131,17 @@ private[connect] class CustomSparkConnectBlockingStub(
}
}
}
+
+ def cloneSession(request: CloneSessionRequest): CloneSessionResponse = {
+ grpcExceptionConverter.convert(
+ request.getSessionId,
+ request.getUserContext,
+ request.getClientType) {
+ retryHandler.retry {
+ stubState.responseValidator.verifyResponse {
+ stub.cloneSession(request)
+ }
+ }
+ }
+ }
}
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
index 48f01a8042a6..3c328681dd9a 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
@@ -29,6 +29,7 @@ import com.google.protobuf.ByteString
import io.grpc._
import org.apache.spark.SparkBuildInfo.{spark_version => SPARK_VERSION}
+import org.apache.spark.annotation.DeveloperApi
import org.apache.spark.connect.proto
import org.apache.spark.connect.proto.UserContext
import org.apache.spark.sql.connect.common.ProtoUtils
@@ -414,6 +415,56 @@ private[sql] class SparkConnectClient(
.build()
artifactManager.cacheArtifact(localRelation.toByteArray)
}
+
+ /**
+ * Clone this client session, creating a new session with the same configuration and shared
+ * state as the current session but with independent runtime state.
+ *
+ * @return
+ * A new SparkConnectClient instance with the cloned session.
+ */
+ @DeveloperApi
+ def cloneSession(): SparkConnectClient = {
+ clone(None)
+ }
+
+ /**
+ * Clone this client session with a custom session ID, creating a new session with the same
+ * configuration and shared state as the current session but with independent runtime state.
+ *
+ * @param newSessionId
+ * Custom session ID to use for the cloned session (must be a valid UUID).
+ * @return
+ * A new SparkConnectClient instance with the cloned session.
+ */
+ @DeveloperApi
+ def clone(newSessionId: String): SparkConnectClient = {
+ clone(Some(newSessionId))
+ }
+
+ private def clone(newSessionId: Option[String]): SparkConnectClient = {
+ val requestBuilder = proto.CloneSessionRequest
+ .newBuilder()
+ .setUserContext(userContext)
+ .setSessionId(sessionId)
+ .setClientType("scala")
+
+ newSessionId.foreach(requestBuilder.setNewSessionId)
+
+ val response: proto.CloneSessionResponse = bstub.cloneSession(requestBuilder.build())
+
+ // Assert that the returned session ID matches the requested ID if one was provided
+ newSessionId.foreach { expectedId =>
+ require(
+ response.getNewSessionId == expectedId,
+ s"Returned session ID '${response.getNewSessionId}' does not match " +
+ s"requested ID '$expectedId'")
+ }
+
+ // Create a new client with the cloned session ID
+ val newConfiguration = configuration.copy(sessionId = Some(response.getNewSessionId))
+ new SparkConnectClient(newConfiguration, configuration.createChannel())
+ }
}
object SparkConnectClient {
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectCloneSessionHandler.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectCloneSessionHandler.scala
new file mode 100644
index 000000000000..c0d6da938932
--- /dev/null
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectCloneSessionHandler.scala
@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.connect.service
+
+import java.util.UUID
+
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+
+class SparkConnectCloneSessionHandler(
+ responseObserver: StreamObserver[proto.CloneSessionResponse])
+ extends Logging {
+
+ def handle(request: proto.CloneSessionRequest): Unit = {
+ val sourceKey = SessionKey(request.getUserContext.getUserId, request.getSessionId)
+ val newSessionId = if (request.hasNewSessionId && request.getNewSessionId.nonEmpty) {
+ request.getNewSessionId
+ } else {
+ UUID.randomUUID().toString
+ }
+ val previouslyObservedSessionId =
+ if (request.hasClientObservedServerSideSessionId &&
+ request.getClientObservedServerSideSessionId.nonEmpty) {
+ Some(request.getClientObservedServerSideSessionId)
+ } else {
+ None
+ }
+ // Get the original session to retrieve its server session ID for validation
+ val originalSessionHolder = SparkConnectService.sessionManager
+ .getIsolatedSession(sourceKey, previouslyObservedSessionId)
+ val clonedSessionHolder = SparkConnectService.sessionManager.cloneSession(
+ sourceKey,
+ newSessionId,
+ previouslyObservedSessionId)
+ val response = proto.CloneSessionResponse
+ .newBuilder()
+ .setSessionId(request.getSessionId)
+ .setNewSessionId(clonedSessionHolder.sessionId)
+ .setServerSideSessionId(originalSessionHolder.serverSessionId)
+ .setNewServerSideSessionId(clonedSessionHolder.serverSessionId)
+ .build()
+ responseObserver.onNext(response)
+ responseObserver.onCompleted()
+ }
+}
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala
index cc1cb95b66c4..13ce2d64256b 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala
@@ -230,6 +230,20 @@ class SparkConnectService(debug: Boolean) extends AsyncService with BindableServ
}
}
+ override def cloneSession(
+ request: proto.CloneSessionRequest,
+ responseObserver: StreamObserver[proto.CloneSessionResponse]): Unit = {
+ try {
+ new SparkConnectCloneSessionHandler(responseObserver).handle(request)
+ } catch {
+ ErrorUtils.handleError(
+ "cloneSession",
+ observer = responseObserver,
+ userId = request.getUserContext.getUserId,
+ sessionId = request.getSessionId)
+ }
+ }
+
private def methodWithCustomMarshallers(
methodDesc: MethodDescriptor[Message, Message]): MethodDescriptor[Message, Message] = {
val recursionLimit =
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala
index 1c3cfd67f132..f28af0379a04 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala
@@ -114,6 +114,73 @@ class SparkConnectSessionManager extends Logging {
Option(getSession(key, None))
}
+ /**
+ * Clone an existing session with a new session ID. Creates a new SessionHolder with a cloned
+ * SparkSession that shares configuration and catalog state but has independent caches and
+ * runtime state.
+ */
+ private[connect] def cloneSession(
+ sourceKey: SessionKey,
+ newSessionId: String,
+ previouslyObservedSessionId: Option[String]): SessionHolder = {
+
+ // Get source session (must exist)
+ val sourceSessionHolder = getIsolatedSession(sourceKey, None)
+
+ previouslyObservedSessionId.foreach(sessionId =>
+ validateSessionId(sourceKey, sourceSessionHolder.session.sessionUUID, sessionId))
+
+ val newKey = SessionKey(sourceKey.userId, newSessionId)
+
+ // Validate new sessionId for clone operation
+ validateCloneTargetSession(newKey)
+
+ // Create cloned session
+ val clonedSessionHolder = getSession(
+ newKey,
+ Some(() => {
+ val session = sessionStore.get(newKey)
+ if (session == null) {
+ // Clone the underlying SparkSession using cloneSession() which preserves
+ // configuration, catalog, session state, temporary views, and registered functions
+ val clonedSparkSession = sourceSessionHolder.session.cloneSession()
+
+ val newHolder = SessionHolder(newKey.userId, newKey.sessionId, clonedSparkSession)
+ newHolder.initializeSession()
+ newHolder
+ } else {
+ // A session was created in the meantime.
+ session
+ }
+ }))
+
+ clonedSessionHolder
+ }
+
+ private def validateCloneTargetSession(newKey: SessionKey): Unit = {
+ // Validate that sessionId is formatted like UUID before creating session.
+ try {
+ UUID.fromString(newKey.sessionId).toString
+ } catch {
+ case _: IllegalArgumentException =>
+ throw new SparkSQLException(
+ errorClass = "INVALID_CLONE_SESSION_REQUEST.TARGET_SESSION_ID_FORMAT",
+ messageParameters = Map("targetSessionId" -> newKey.sessionId))
+ }
+ // Validate that session with that key has not been already closed.
+ if (closedSessionsCache.getIfPresent(newKey) != null) {
+ throw new SparkSQLException(
+ errorClass = "INVALID_CLONE_SESSION_REQUEST.TARGET_SESSION_ID_ALREADY_CLOSED",
+ messageParameters = Map("targetSessionId" -> newKey.sessionId))
+ }
+ // Validate that session with that key does not already exist.
+ if (sessionStore.containsKey(newKey)) {
+ throw new SparkSQLException(
+ errorClass = "INVALID_CLONE_SESSION_REQUEST.TARGET_SESSION_ID_ALREADY_EXISTS",
+ messageParameters = Map("targetSessionId" -> newKey.sessionId))
+ }
+ }
+
private def getSession(key: SessionKey, default: Option[() => SessionHolder]): SessionHolder = {
schedulePeriodicChecks() // Starts the maintenance thread if it hasn't started yet.
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectCloneSessionSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectCloneSessionSuite.scala
new file mode 100644
index 000000000000..922c239526f3
--- /dev/null
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectCloneSessionSuite.scala
@@ -0,0 +1,279 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.connect.service
+
+import java.util.UUID
+
+import org.scalatest.BeforeAndAfterEach
+
+import org.apache.spark.SparkSQLException
+import org.apache.spark.sql.test.SharedSparkSession
+
+class SparkConnectCloneSessionSuite extends SharedSparkSession with BeforeAndAfterEach {
+
+ override def beforeEach(): Unit = {
+ super.beforeEach()
+ SparkConnectService.sessionManager.invalidateAllSessions()
+ }
+
+ test("clone session with invalid target session ID format") {
+ val sourceKey = SessionKey("testUser", UUID.randomUUID.toString)
+ val invalidSessionId = "not-a-valid-uuid"
+
+ // Create the source session first
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(sourceKey, None)
+
+ val ex = intercept[SparkSQLException] {
+ SparkConnectService.sessionManager.cloneSession(sourceKey, invalidSessionId, None)
+ }
+
+ assert(ex.getCondition == "INVALID_CLONE_SESSION_REQUEST.TARGET_SESSION_ID_FORMAT")
+ assert(ex.getMessage.contains("Target session ID not-a-valid-uuid"))
+ assert(ex.getMessage.contains("must be an UUID string"))
+ }
+
+ test("clone session with target session ID already closed") {
+ val sourceKey = SessionKey("testUser", UUID.randomUUID.toString)
+ val targetSessionId = UUID.randomUUID.toString
+ val targetKey = SessionKey("testUser", targetSessionId)
+
+ // Create and then close a session to put it in the closed cache
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(sourceKey, None)
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(targetKey, None)
+ SparkConnectService.sessionManager.closeSession(targetKey)
+
+ val ex = intercept[SparkSQLException] {
+ SparkConnectService.sessionManager.cloneSession(sourceKey, targetSessionId, None)
+ }
+
+ assert(
+ ex.getCondition ==
+ "INVALID_CLONE_SESSION_REQUEST.TARGET_SESSION_ID_ALREADY_CLOSED")
+ assert(ex.getMessage.contains(s"target session ID $targetSessionId"))
+ assert(ex.getMessage.contains("was previously closed"))
+ }
+
+ test("clone session with target session ID already exists") {
+ val sourceKey = SessionKey("testUser", UUID.randomUUID.toString)
+ val targetSessionId = UUID.randomUUID.toString
+ val targetKey = SessionKey("testUser", targetSessionId)
+
+ // Create both source and target sessions
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(sourceKey, None)
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(targetKey, None)
+
+ val ex = intercept[SparkSQLException] {
+ SparkConnectService.sessionManager.cloneSession(sourceKey, targetSessionId, None)
+ }
+
+ assert(
+ ex.getCondition ==
+ "INVALID_CLONE_SESSION_REQUEST.TARGET_SESSION_ID_ALREADY_EXISTS")
+ assert(ex.getMessage.contains(s"target session ID $targetSessionId"))
+ assert(ex.getMessage.contains("already exists"))
+ }
+
+ test("clone session with source session not found") {
+ val sourceKey = SessionKey("testUser", UUID.randomUUID.toString)
+ val targetSessionId = UUID.randomUUID.toString
+
+ // Don't create the source session, so it doesn't exist
+ val ex = intercept[SparkSQLException] {
+ SparkConnectService.sessionManager.cloneSession(sourceKey, targetSessionId, None)
+ }
+
+ // Source session errors should remain as standard INVALID_HANDLE errors
+ assert(ex.getCondition == "INVALID_HANDLE.SESSION_NOT_FOUND")
+ assert(ex.getMessage.contains("Session not found"))
+ }
+
+ test("successful clone session creates new session") {
+ val sourceKey = SessionKey("testUser", UUID.randomUUID.toString)
+ val targetSessionId = UUID.randomUUID.toString
+
+ // Create source session
+ val sourceSession = SparkConnectService.sessionManager
+ .getOrCreateIsolatedSession(sourceKey, None)
+
+ // Clone the session
+ val clonedSession = SparkConnectService.sessionManager
+ .cloneSession(sourceKey, targetSessionId, None)
+
+ // Verify the cloned session has the expected session ID
+ assert(clonedSession.sessionId == targetSessionId)
+ assert(clonedSession.sessionId != sourceSession.sessionId)
+
+ // Both sessions should be different objects
+ assert(clonedSession != sourceSession)
+ }
+
+ test("cloned session copies all state (configs, temp views, UDFs, current database)") {
+ val sourceKey = SessionKey("testUser", UUID.randomUUID.toString)
+ val targetSessionId = UUID.randomUUID.toString
+
+ // Create source session and set up state
+ val sourceSession = SparkConnectService.sessionManager
+ .getOrCreateIsolatedSession(sourceKey, None)
+
+ // Set SQL configs
+ sourceSession.session.conf.set("spark.sql.custom.test.config", "test-value")
+ sourceSession.session.conf.set("spark.sql.shuffle.partitions", "42")
+
+ // Create temp views
+ sourceSession.session.sql("CREATE TEMPORARY VIEW temp_view1 AS SELECT 1 as col1, 2 as col2")
+ sourceSession.session.sql("CREATE TEMPORARY VIEW temp_view2 AS SELECT 3 as col3")
+
+ // Register UDFs
+ sourceSession.session.udf.register("double_value", (x: Int) => x * 2)
+ sourceSession.session.udf.register("concat_strings", (a: String, b: String) => a + b)
+
+ // Change database
+ sourceSession.session.sql("CREATE DATABASE IF NOT EXISTS test_db")
+ sourceSession.session.sql("USE test_db")
+
+ // Clone the session
+ val clonedSession = SparkConnectService.sessionManager
+ .cloneSession(sourceKey, targetSessionId, None)
+
+ // Verify all state is copied
+
+ // Configs are copied
+ assert(clonedSession.session.conf.get("spark.sql.custom.test.config") == "test-value")
+ assert(clonedSession.session.conf.get("spark.sql.shuffle.partitions") == "42")
+
+ // Temp views are accessible
+ val view1Result = clonedSession.session.sql("SELECT * FROM temp_view1").collect()
+ assert(view1Result.length == 1)
+ assert(view1Result(0).getInt(0) == 1)
+ assert(view1Result(0).getInt(1) == 2)
+
+ val view2Result = clonedSession.session.sql("SELECT * FROM temp_view2").collect()
+ assert(view2Result.length == 1)
+ assert(view2Result(0).getInt(0) == 3)
+
+ // UDFs are accessible
+ val udfResult1 = clonedSession.session.sql("SELECT double_value(5)").collect()
+ assert(udfResult1(0).getInt(0) == 10)
+
+ val udfResult2 =
+ clonedSession.session.sql("SELECT concat_strings('hello', 'world')").collect()
+ assert(udfResult2(0).getString(0) == "helloworld")
+
+ // Current database is copied
+ assert(clonedSession.session.catalog.currentDatabase == "test_db")
+ }
+
+ test("sessions are independent after cloning (configs, temp views, UDFs)") {
+ val sourceKey = SessionKey("testUser", UUID.randomUUID.toString)
+ val targetSessionId = UUID.randomUUID.toString
+
+ // Create and set up source session
+ val sourceSession = SparkConnectService.sessionManager
+ .getOrCreateIsolatedSession(sourceKey, None)
+ sourceSession.session.conf.set("spark.sql.custom.config", "initial")
+ sourceSession.session.sql("CREATE TEMPORARY VIEW shared_view AS SELECT 1 as value")
+ sourceSession.session.udf.register("shared_udf", (x: Int) => x + 1)
+
+ // Clone the session
+ val clonedSession = SparkConnectService.sessionManager
+ .cloneSession(sourceKey, targetSessionId, None)
+
+ // Test independence of configs
+ sourceSession.session.conf.set("spark.sql.custom.config", "modified-source")
+ clonedSession.session.conf.set("spark.sql.custom.config", "modified-clone")
+ assert(sourceSession.session.conf.get("spark.sql.custom.config") == "modified-source")
+ assert(clonedSession.session.conf.get("spark.sql.custom.config") == "modified-clone")
+
+ // Test independence of temp views - modify shared view differently
+ sourceSession.session.sql(
+ "CREATE OR REPLACE TEMPORARY VIEW shared_view AS SELECT 10 as value")
+ clonedSession.session.sql(
+ "CREATE OR REPLACE TEMPORARY VIEW shared_view AS SELECT 20 as value")
+
+ // Each session should see its own version of the view
+ val sourceViewResult = sourceSession.session.sql("SELECT * FROM shared_view").collect()
+ assert(sourceViewResult(0).getInt(0) == 10)
+
+ val cloneViewResult = clonedSession.session.sql("SELECT * FROM shared_view").collect()
+ assert(cloneViewResult(0).getInt(0) == 20)
+
+ // Test independence of UDFs
+ sourceSession.session.udf.register("shared_udf", (x: Int) => x + 10)
+ clonedSession.session.udf.register("shared_udf", (x: Int) => x + 100)
+ assert(sourceSession.session.sql("SELECT shared_udf(5)").collect()(0).getInt(0) == 15)
+ assert(clonedSession.session.sql("SELECT shared_udf(5)").collect()(0).getInt(0) == 105)
+ }
+
+ test("cloned session copies artifacts and maintains independence") {
+ import java.nio.charset.StandardCharsets
+ import java.nio.file.{Files, Paths}
+
+ val sourceKey = SessionKey("testUser", UUID.randomUUID.toString)
+ val targetSessionId = UUID.randomUUID.toString
+
+ // Create source session
+ val sourceSession = SparkConnectService.sessionManager
+ .getOrCreateIsolatedSession(sourceKey, None)
+
+ // Add some test artifacts to source session
+ val tempFile = Files.createTempFile("test-artifact", ".txt")
+ Files.write(tempFile, "test content".getBytes(StandardCharsets.UTF_8))
+
+ // Add artifact to source session
+ val remotePath = Paths.get("test/artifact.txt")
+ sourceSession.artifactManager.addArtifact(remotePath, tempFile, None)
+
+ // Clone the session
+ val clonedSession = SparkConnectService.sessionManager
+ .cloneSession(sourceKey, targetSessionId, None)
+
+ // Verify sessions have different artifact managers
+ assert(sourceSession.artifactManager ne clonedSession.artifactManager)
+ assert(sourceSession.session.sessionUUID != clonedSession.session.sessionUUID)
+
+ // Test independence: add new artifacts to each session
+ val sourceOnlyFile = Files.createTempFile("source-only", ".txt")
+ Files.write(sourceOnlyFile, "source only content".getBytes(StandardCharsets.UTF_8))
+ val sourceOnlyPath = Paths.get("jars/source.jar")
+ sourceSession.artifactManager.addArtifact(sourceOnlyPath, sourceOnlyFile, None)
+
+ val clonedOnlyFile = Files.createTempFile("cloned-only", ".txt")
+ Files.write(clonedOnlyFile, "cloned only content".getBytes(StandardCharsets.UTF_8))
+ val clonedOnlyPath = Paths.get("jars/cloned.jar")
+ clonedSession.artifactManager.addArtifact(clonedOnlyPath, clonedOnlyFile, None)
+
+ // Use getAddedJars to verify independence (since it's a public API)
+ val sourceJars = sourceSession.artifactManager.getAddedJars.map(_.toString)
+ val clonedJars = clonedSession.artifactManager.getAddedJars.map(_.toString)
+
+ // Source should have source.jar but not cloned.jar
+ assert(sourceJars.exists(_.contains("source.jar")))
+ assert(!sourceJars.exists(_.contains("cloned.jar")))
+
+ // Cloned should have cloned.jar but not source.jar
+ assert(clonedJars.exists(_.contains("cloned.jar")))
+ assert(!clonedJars.exists(_.contains("source.jar")))
+
+ // Clean up
+ Files.deleteIfExists(tempFile)
+ Files.deleteIfExists(sourceOnlyFile)
+ Files.deleteIfExists(clonedOnlyFile)
+ sourceSession.artifactManager.close()
+ clonedSession.artifactManager.close()
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org