fix: Remove 'Dataproc Session' link from the cell output and update 'View Session Details' button url. (#163)

Author: bansalaayushi23

google/cloud/dataproc_spark_connect/session.py

@@ -67,6 +67,10 @@
     "goog-colab-notebook-id",
 }
 
+_DATAPROC_SESSIONS_BASE_URL = (
+    "https://console.cloud.google.com/dataproc/interactive"
+)
+
 
 def _is_valid_label_value(value: str) -> bool:
     """
@@ -494,15 +498,21 @@ def _wait_for_session_available(
             )
 
         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}"
+            session_url = f"{_DATAPROC_SESSIONS_BASE_URL}/{self._region}/{session_id}?project={self._project_id}"
             plain_message = f"Creating Dataproc Session: {session_url}"
-            html_element = f"""
+            if environment.is_colab_enterprise():
+                html_element = f"""
                 <div>
                     <p>Creating Dataproc Spark Session<p>
-                    <p><a href="{session_url}">Dataproc Session</a></p>
                 </div>
-            """
-
+                """
+            else:
+                html_element = f"""
+                    <div>
+                        <p>Creating Dataproc Spark Session<p>
+                        <p><a href="{session_url}">Dataproc Session</a></p>
+                    </div>
+                """
             self._output_element_or_message(plain_message, html_element)
 
         def _print_session_created_message(self):
@@ -554,7 +564,7 @@ def _get_exiting_active_session(
 
             if session_response is not None:
                 print(
-                    f"Using existing Dataproc Session (configuration changes may not be applied): https://console.cloud.google.com/dataproc/interactive/{self._region}/{s8s_session_id}?project={self._project_id}"
+                    f"Using existing Dataproc Session (configuration changes may not be applied): {_DATAPROC_SESSIONS_BASE_URL}/{self._region}/{s8s_session_id}?project={self._project_id}"
                 )
                 self._display_view_session_details_button(s8s_session_id)
                 if session is None:
@@ -798,7 +808,7 @@ def _display_view_session_details_button(self, session_id):
                 return
 
             try:
-                session_url = f"https://console.cloud.google.com/dataproc/interactive/sessions/{session_id}/locations/{self._region}?project={self._project_id}"
+                session_url = f"{_DATAPROC_SESSIONS_BASE_URL}/{self._region}/{session_id}?project={self._project_id}"
                 from IPython.core.interactiveshell import InteractiveShell
 
                 if not InteractiveShell.initialized():
@@ -1060,7 +1070,7 @@ def _repr_html_(self) -> str:
             <div>No Active Dataproc Session</div>
             """
 
-        s8s_session = f"https://console.cloud.google.com/dataproc/interactive/{self._region}/{self._active_s8s_session_id}"
+        s8s_session = f"{_DATAPROC_SESSIONS_BASE_URL}/{self._region}/{self._active_s8s_session_id}"
         ui = f"{s8s_session}/sparkApplications/applications"
         return f"""
         <div>
@@ -1087,7 +1097,7 @@ def _display_operation_link(self, operation_id: str):
         )
 
         url = (
-            f"https://console.cloud.google.com/dataproc/interactive/{self._region}/"
+            f"{_DATAPROC_SESSIONS_BASE_URL}/{self._region}/"
             f"{self._active_s8s_session_id}/sparkApplications/application;"
             f"associatedSqlOperationId={operation_id}?project={self._project_id}"
         )

tests/unit/test_session.py

@@ -38,6 +38,10 @@
 from pyspark.sql.connect.proto import Command, ConfigResponse, ExecutePlanRequest, Plan, Relation, SQL, SqlCommand, UserContext
 from unittest import mock
 
+_DATAPROC_SESSIONS_BASE_URL = (
+    "https://console.cloud.google.com/dataproc/interactive"
+)
+
 
 class DataprocRemoteSparkSessionBuilderTests(unittest.TestCase):
 
@@ -163,7 +167,7 @@ def test_create_spark_session_with_default_notebook_behavior(
         mock_ipython_utils = mock.sys.modules[
             "google.cloud.aiplatform.utils"
         ]._ipython_utils
-        test_session_url = f"https://console.cloud.google.com/dataproc/interactive/sessions/{session_id}/locations/test-region?project=test-project"
+        test_session_url = f"{_DATAPROC_SESSIONS_BASE_URL}/test-region/{session_id}?project=test-project"
         mock_display_link = mock_ipython_utils.display_link
         mock.patch.dict(
             os.environ,
@@ -1189,7 +1193,7 @@ def test_display_button_with_aiplatform_installed_ipython_interactive(
         mock_ipython_utils = mock.sys.modules[
             "google.cloud.aiplatform.utils"
         ]._ipython_utils
-        test_session_url = "https://console.cloud.google.com/dataproc/interactive/sessions/test_session/locations/test-region?project=test-project"
+        test_session_url = f"{_DATAPROC_SESSIONS_BASE_URL}/test-region/test_session?project=test-project"
 
         mock_display_link = mock_ipython_utils.display_link
         DataprocSparkSession.builder._display_view_session_details_button(
@@ -1230,6 +1234,56 @@ def test_display_button_with_aiplatform_installed_ipython_non_interactive(
         )
         mock_display_link.assert_not_called()
 
+    @mock.patch(
+        "IPython.core.interactiveshell.InteractiveShell.initialized",
+        return_value=False,
+    )
+    @mock.patch("IPython.display.display")
+    def test_display_session_link_on_creation_colab_enterprise(
+        self,
+        mock_display,
+        _mock_ipy,
+    ):
+        mock.patch.dict(
+            os.environ,
+            {
+                "VERTEX_PRODUCT": "COLAB_ENTERPRISE",
+            },
+        ).start()
+        DataprocSparkSession.builder._display_session_link_on_creation(
+            "test_session"
+        )
+
+        mock_display.assert_called_once()
+        args, _ = mock_display.call_args
+        html_output = args[0].data
+        self.assertIn("Creating Dataproc Spark Session", html_output)
+        self.assertNotIn("Dataproc Session", html_output)
+
+    @mock.patch(
+        "IPython.core.interactiveshell.InteractiveShell.initialized",
+        return_value=False,
+    )
+    @mock.patch("IPython.display.display")
+    def test_display_session_link_on_creation_not_colab_enterprise(
+        self,
+        mock_display,
+        _mock_ipy,
+    ):
+        mock.patch.dict(
+            os.environ,
+            {},
+        ).start()
+        DataprocSparkSession.builder._display_session_link_on_creation(
+            "test_session"
+        )
+
+        mock_display.assert_called_once()
+        args, _ = mock_display.call_args
+        html_output = args[0].data
+        self.assertIn("Creating Dataproc Spark Session", html_output)
+        self.assertIn("Dataproc Session", html_output)
+
     def test_is_valid_label_value(self):
         # Valid label values
         self.assertTrue(_is_valid_label_value("valid-label-123"))