Use vscode:// deep links for session URLs in VS Code

google/cloud/managed_spark_connect/session.py

@@ -71,6 +71,18 @@
     "https://console.cloud.google.com/dataproc/interactive"
 )
 
+_VSCODE_SESSION_URI_BASE = (
+    "vscode://googlecloudtools.datacloud/dataproc/sessions"
+)
+
+
+def _build_session_details_url(
+    region: Optional[str], project_id: Optional[str], session_id: str
+) -> str:
+    if environment.is_vscode():
+        return f"{_VSCODE_SESSION_URI_BASE}/{session_id}?project={project_id}&location={region}"
+    return f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{region}/{session_id}?project={project_id}"
+
 
 def _is_valid_label_value(value: str) -> bool:
     """
@@ -506,7 +518,9 @@ def _wait_for_session_available(
             )
 
         def _display_session_link_on_creation(self, session_id):
-            session_url = f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/{session_id}?project={self._project_id}"
+            session_url = _build_session_details_url(
+                self._region, self._project_id, session_id
+            )
             plain_message = (
                 f"Creating Managed Spark Connect Session: {session_url}"
             )
@@ -573,8 +587,11 @@ def _get_exiting_active_session(
                 session = ManagedSparkSession._default_session
 
             if session_response is not None:
+                session_url = _build_session_details_url(
+                    self._region, self._project_id, s8s_session_id
+                )
                 print(
-                    f"Using existing Managed Spark Session (configuration changes may not be applied): {_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/{s8s_session_id}?project={self._project_id}"
+                    f"Using existing Managed Spark Session (configuration changes may not be applied): {session_url}"
                 )
                 self._display_view_session_details_button(s8s_session_id)
                 if session is None:
@@ -1107,14 +1124,20 @@ def _repr_html_(self) -> str:
             <div>No Active Managed Spark Session</div>
             """
 
-        s8s_session = f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/{self._active_s8s_session_id}"
-        ui = f"{s8s_session}/sparkApplications/applications"
+        session_url = _build_session_details_url(
+            self._region, self._project_id, self._active_s8s_session_id
+        )
+        ui = (
+            f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/"
+            f"{self._active_s8s_session_id}/sparkApplications/applications"
+            f"?project={self._project_id}"
+        )
         return f"""
         <div>
             <p><b>Spark Connect</b></p>
 
-            <p><a href="{s8s_session}?project={self._project_id}">Managed Spark Session</a></p>
-            <p><a href="{ui}?project={self._project_id}">Spark UI</a></p>
+            <p><a href="{session_url}">Managed Spark Session</a></p>
+            <p><a href="{ui}">Spark UI</a></p>
         </div>
         """
 

tests/unit/test_session.py

@@ -1259,6 +1259,103 @@ def test_display_session_link_on_creation_not_colab_enterprise(
         self.assertIn("Creating Managed Spark Connect Session", html_output)
         self.assertIn("Managed Spark Session", html_output)
 
+    @mock.patch(
+        "IPython.core.interactiveshell.InteractiveShell.initialized",
+        return_value=True,
+    )
+    @mock.patch("IPython.display.display")
+    def test_display_session_link_on_creation_vscode(
+        self,
+        mock_display,
+        _mock_ipy,
+    ):
+        mock.patch.dict(
+            os.environ,
+            {
+                "VSCODE_PID": "12345",
+            },
+        ).start()
+        ManagedSparkSession.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 Managed Spark Connect Session", html_output)
+        self.assertIn(
+            "vscode://googlecloudtools.datacloud/dataproc/sessions/"
+            "test_session?project=test-project&location=test-region",
+            html_output,
+        )
+        self.assertNotIn(_MANAGED_SPARK_SESSIONS_BASE_URL, html_output)
+
+    @mock.patch.object(ManagedSparkSession, "getActiveSession")
+    @mock.patch(
+        "google.cloud.managed_spark_connect.session.get_active_s8s_session_response"
+    )
+    def test_get_exiting_active_session_prints_vscode_url(
+        self,
+        mock_get_response,
+        mock_get_active_session,
+    ):
+        mock.patch.dict(
+            os.environ,
+            {
+                "VSCODE_PID": "12345",
+            },
+        ).start()
+        mock_get_response.return_value = mock.Mock()
+        mock_get_active_session.return_value = mock.Mock()
+        ManagedSparkSession._active_s8s_session_id = "test_session"
+        self.addCleanup(
+            setattr, ManagedSparkSession, "_active_s8s_session_id", None
+        )
+
+        with mock.patch("builtins.print") as mock_print:
+            ManagedSparkSession.builder._get_exiting_active_session()
+
+        printed = "\n".join(
+            str(call.args[0]) for call in mock_print.call_args_list
+        )
+        self.assertIn(
+            "vscode://googlecloudtools.datacloud/dataproc/sessions/"
+            "test_session?project=test-project&location=test-region",
+            printed,
+        )
+        self.assertNotIn(_MANAGED_SPARK_SESSIONS_BASE_URL, printed)
+
+    def test_repr_html_uses_vscode_url_for_session_link(self):
+        mock.patch.dict(
+            os.environ,
+            {
+                "VSCODE_PID": "12345",
+            },
+        ).start()
+        ManagedSparkSession._project_id = "test-project"
+        ManagedSparkSession._region = "test-region"
+        ManagedSparkSession._active_s8s_session_id = "test_session"
+        self.addCleanup(setattr, ManagedSparkSession, "_project_id", None)
+        self.addCleanup(setattr, ManagedSparkSession, "_region", None)
+        self.addCleanup(
+            setattr, ManagedSparkSession, "_active_s8s_session_id", None
+        )
+
+        html = object.__new__(ManagedSparkSession)._repr_html_()
+
+        self.assertIn(
+            '<a href="vscode://googlecloudtools.datacloud/dataproc/sessions/'
+            'test_session?project=test-project&location=test-region">'
+            "Managed Spark Session</a>",
+            html,
+        )
+        self.assertIn(
+            f'<a href="{_MANAGED_SPARK_SESSIONS_BASE_URL}/test-region/'
+            'test_session/sparkApplications/applications?project=test-project">'
+            "Spark UI</a>",
+            html,
+        )
+
     def test_is_valid_label_value(self):
         # Valid label values
         self.assertTrue(_is_valid_label_value("valid-label-123"))