fix: pysparkvalue error

Author: zhiweilin-svg

google/cloud/dataproc_spark_connect/session.py

@@ -472,6 +472,27 @@ def create_session_pbar():
                     session_response, dataproc_config.name
                 )
 
+        def _wait_for_session_available(
+            self, session_name: str, timeout: int = 300
+        ) -> Session:
+            start_time = time.time()
+            while time.time() - start_time < timeout:
+                try:
+                    session = self.session_controller_client.get_session(
+                        name=session_name
+                    )
+                    if "Spark Connect Server" in session.runtime_info.endpoints:
+                        return session
+                    time.sleep(5)
+                except Exception as e:
+                    logger.warning(
+                        f"Error while polling for Spark Connect endpoint: {e}"
+                    )
+                    time.sleep(5)
+            raise RuntimeError(
+                f"Spark Connect endpoint not available for session {session_name} after {timeout} seconds."
+            )
+
         def _display_session_link_on_creation(self, session_id):
             session_url = f"https://console.cloud.google.com/dataproc/interactive/{self._region}/{session_id}?project={self._project_id}"
             plain_message = f"Creating Dataproc Session: {session_url}"
@@ -537,6 +558,9 @@ def _get_exiting_active_session(
                 )
                 self._display_view_session_details_button(s8s_session_id)
                 if session is None:
+                    session_response = self._wait_for_session_available(
+                        session_name
+                    )
                     session = self.__create_spark_connect_session_from_s8s(
                         session_response, session_name
                     )

tests/unit/test_session.py

@@ -2393,6 +2393,69 @@ def test_execution_progress_handler(
             )
             self.stopSession(mock_session_controller_client_instance, session)
 
+    @mock.patch("time.sleep", return_value=None)
+    @mock.patch("google.cloud.dataproc_v1.SessionControllerClient")
+    def test_wait_for_session_available_success(
+        self, mock_session_controller_client, mock_sleep
+    ):
+        """Test that the method waits and returns the session when the endpoint appears."""
+        mock_client = mock_session_controller_client.return_value
+        session_name = "projects/test-project/locations/test-region/sessions/test-session"
+
+        # Session without the endpoint
+        session_pending = Session()
+        session_pending.name = session_name
+
+        # Session with the endpoint
+        session_ready = Session()
+        session_ready.name = session_name
+        session_ready.runtime_info.endpoints[
+            "Spark Connect Server"
+        ] = "sc://example.com:443"
+
+        # Mock get_session to return pending, then ready
+        mock_client.get_session.side_effect = [
+            session_pending,
+            session_pending,
+            session_ready,
+        ]
+
+        builder = DataprocSparkSession.Builder()
+        builder._session_controller_client = mock_client  # Inject the mock client
+
+        result = builder._wait_for_session_available(session_name, timeout=10)
+
+        self.assertEqual(result, session_ready)
+        self.assertEqual(mock_client.get_session.call_count, 3)
+        self.assertEqual(mock_sleep.call_count, 2)
+
+    @mock.patch("time.sleep", return_value=None)
+    @mock.patch("google.cloud.dataproc_v1.SessionControllerClient")
+    def test_wait_for_session_available_timeout(
+        self, mock_session_controller_client, mock_sleep
+    ):
+        """Test that the method raises RuntimeError on timeout."""
+        mock_client = mock_session_controller_client.return_value
+        session_name = "projects/test-project/locations/test-region/sessions/test-session"
+
+        # Session that never gets the endpoint
+        session_pending = Session()
+        session_pending.name = session_name
+
+        mock_client.get_session.return_value = session_pending
+
+        builder = DataprocSparkSession.Builder()
+        builder._session_controller_client = mock_client  # Inject the mock client
+
+        with self.assertRaises(RuntimeError) as context:
+            # Use a short timeout for the test
+            builder._wait_for_session_available(session_name, timeout=1)
+
+        self.assertIn(
+            f"Spark Connect endpoint not available for session {session_name}",
+            str(context.exception),
+        )
+
 
 class SessionIdValidationTests(unittest.TestCase):
     """Test cases for session ID validation and custom session ID functionality."""