feat: support sharable session termination by adding terminate option to stop() function (#162)

Author: fangyh20

google/cloud/dataproc_spark_connect/session.py

@@ -1179,27 +1179,64 @@ def addArtifacts(
     def _get_active_session_file_path():
         return os.getenv("DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH")
 
-    def stop(self) -> None:
+    def stop(self, terminate: Optional[bool] = None) -> None:
+        """
+        Stop the Spark session and optionally terminate the server-side session.
+
+        Parameters
+        ----------
+        terminate : bool, optional
+            Control server-side termination behavior.
+
+            - None (default): Auto-detect based on session type
+
+              - Managed sessions (auto-generated ID): terminate server
+              - Named sessions (custom ID): client-side cleanup only
+
+            - True: Always terminate the server-side session
+            - False: Never terminate the server-side session (client cleanup only)
+
+        Examples
+        --------
+        Auto-detect termination behavior (existing behavior):
+
+        >>> spark.stop()
+
+        Force terminate a named session:
+
+        >>> spark.stop(terminate=True)
+
+        Prevent termination of a managed session:
+
+        >>> spark.stop(terminate=False)
+        """
         with DataprocSparkSession._lock:
             if DataprocSparkSession._active_s8s_session_id is not None:
-                # Check if this is a managed session (auto-generated ID) or unmanaged session (custom ID)
-                if DataprocSparkSession._active_session_uses_custom_id:
-                    # Unmanaged session (custom ID): Only clean up client-side state
-                    # Don't terminate as it might be in use by other notebooks or clients
-                    logger.debug(
-                        f"Stopping unmanaged session {DataprocSparkSession._active_s8s_session_id} without termination"
+                # Determine if we should terminate the server-side session
+                if terminate is None:
+                    # Auto-detect: managed sessions terminate, named sessions don't
+                    should_terminate = (
+                        not DataprocSparkSession._active_session_uses_custom_id
                     )
                 else:
-                    # Managed session (auto-generated ID): Use original behavior and terminate
+                    should_terminate = terminate
+
+                if should_terminate:
+                    # Terminate the server-side session
                     logger.debug(
-                        f"Terminating managed session {DataprocSparkSession._active_s8s_session_id}"
+                        f"Terminating session {DataprocSparkSession._active_s8s_session_id}"
                     )
                     terminate_s8s_session(
                         DataprocSparkSession._project_id,
                         DataprocSparkSession._region,
                         DataprocSparkSession._active_s8s_session_id,
                         self._client_options,
                     )
+                else:
+                    # Client-side cleanup only
+                    logger.debug(
+                        f"Stopping session {DataprocSparkSession._active_s8s_session_id} without termination"
+                    )
 
                 self._remove_stopped_session_from_file()
 

tests/integration/test_session.py

@@ -666,6 +666,96 @@ def test_sparksql_magic_with_dataproc_session(connect_session):
     assert row["joined_string"] == "Dataproc-Spark"
 
 
+def test_stop_named_session_with_terminate_true(
+    auth_type,
+    test_project,
+    test_region,
+    session_controller_client,
+    os_environment,
+):
+    """Test that stop(terminate=True) terminates a named session on the server."""
+    # Use a randomized session ID to avoid conflicts
+    custom_session_id = f"test-terminate-true-{uuid.uuid4().hex[:8]}"
+
+    # Create a session with custom ID
+    spark = (
+        DataprocSparkSession.builder.dataprocSessionId(custom_session_id)
+        .projectId(test_project)
+        .location(test_region)
+        .getOrCreate()
+    )
+
+    # Verify session is created
+    assert DataprocSparkSession._active_s8s_session_id == custom_session_id
+    session_name = f"projects/{test_project}/locations/{test_region}/sessions/{custom_session_id}"
+
+    # Test basic functionality
+    df = spark.createDataFrame([(1, "test")], ["id", "value"])
+    assert df.count() == 1
+
+    # Stop with terminate=True
+    spark.stop(terminate=True)
+
+    # Verify client-side cleanup
+    assert DataprocSparkSession._active_s8s_session_id is None
+
+    # Verify server-side session is terminating or terminated
+    get_session_request = GetSessionRequest()
+    get_session_request.name = session_name
+    session = session_controller_client.get_session(get_session_request)
+
+    assert session.state in [
+        Session.State.TERMINATING,
+        Session.State.TERMINATED,
+    ]
+
+
+def test_stop_managed_session_with_terminate_false(
+    auth_type,
+    test_project,
+    test_region,
+    session_controller_client,
+    os_environment,
+):
+    """Test that stop(terminate=False) does NOT terminate a managed session on the server."""
+    # Create a managed session (auto-generated ID)
+    spark = (
+        DataprocSparkSession.builder.projectId(test_project)
+        .location(test_region)
+        .getOrCreate()
+    )
+
+    # Verify it's a managed session (auto-generated ID)
+    assert DataprocSparkSession._active_s8s_session_id is not None
+    assert DataprocSparkSession._active_session_uses_custom_id is False
+    session_id = DataprocSparkSession._active_s8s_session_id
+    session_name = (
+        f"projects/{test_project}/locations/{test_region}/sessions/{session_id}"
+    )
+
+    # Test basic functionality
+    df = spark.createDataFrame([(1, "test")], ["id", "value"])
+    assert df.count() == 1
+
+    # Stop with terminate=False (prevent server-side termination)
+    spark.stop(terminate=False)
+
+    # Verify client-side cleanup
+    assert DataprocSparkSession._active_s8s_session_id is None
+
+    # Verify server-side session is still ACTIVE (not terminated)
+    get_session_request = GetSessionRequest()
+    get_session_request.name = session_name
+    session = session_controller_client.get_session(get_session_request)
+
+    assert session.state == Session.State.ACTIVE
+
+    # Clean up: terminate the session manually
+    terminate_session_request = TerminateSessionRequest()
+    terminate_session_request.name = session_name
+    session_controller_client.terminate_session(terminate_session_request)
+
+
 @pytest.fixture
 def batch_workload_env(monkeypatch):
     """Sets DATAPROC_WORKLOAD_TYPE to 'batch' for a test."""