diff --git a/google/cloud/dataproc_spark_connect/session.py b/google/cloud/dataproc_spark_connect/session.py index 09b9f62..d2fce92 100644 --- a/google/cloud/dataproc_spark_connect/session.py +++ b/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 ) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index c196dfe..4e8977e 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -1143,7 +1143,7 @@ def test_create_session_with_default_datasource_env_var( mock_session_controller_client_instance.create_session.reset_mock() mock_logger.warning.reset_mock() - @mock.patch("sys.modules", {"google.cloud.aiplatform": None}) + @mock.patch.dict("sys.modules", {"google.cloud.aiplatform": None}) @mock.patch( "IPython.core.interactiveshell.InteractiveShell.initialized", return_value=True, @@ -2393,6 +2393,77 @@ 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."""