Spark Connect
- + """ From 728f3419f00c9b37d424d56f3c45ab0d4c4720e5 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:47:57 -0700 Subject: [PATCH 06/10] fix: fall back to console url if the vscode extension isn't installed The vscode:// deep link only resolves if googlecloudtools.datacloud is installed; otherwise VS Code shows a generic, confusing error. Check via `code --list-extensions` and fall back to the Cloud Console url when it's not present or the check can't run. --- .../managed_spark_connect/environment.py | 30 +++++++++++ google/cloud/managed_spark_connect/session.py | 9 ++-- tests/unit/test_environment.py | 38 +++++++++++++ tests/unit/test_session.py | 53 ++++++++++++++++++- 4 files changed, 125 insertions(+), 5 deletions(-) diff --git a/google/cloud/managed_spark_connect/environment.py b/google/cloud/managed_spark_connect/environment.py index e19dd97..bfb8135 100644 --- a/google/cloud/managed_spark_connect/environment.py +++ b/google/cloud/managed_spark_connect/environment.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import os +import subprocess import sys from typing import Callable, Tuple, List @@ -27,6 +29,34 @@ def is_vscode() -> bool: return os.getenv("VSCODE_PID") is not None +@functools.lru_cache(maxsize=1) +def _installed_vscode_extensions() -> frozenset: + try: + result = subprocess.run( + ["code", "--list-extensions"], + capture_output=True, + text=True, + timeout=5, + check=True, + ) + except (OSError, subprocess.SubprocessError): + return frozenset() + return frozenset( + line.strip().lower() + for line in result.stdout.splitlines() + if line.strip() + ) + + +def is_vscode_extension_installed(extension_id: str) -> bool: + """True if the given VS Code extension id is installed. + + Fails closed (returns False) if the `code` CLI is unavailable, times + out, or exits non-zero. + """ + return extension_id.lower() in _installed_vscode_extensions() + + def is_jupyter() -> bool: """True if running in a Jupyter environment.""" return os.getenv("JPY_PARENT_PID") is not None diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index bb5d813..258605a 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -71,15 +71,16 @@ "https://console.cloud.google.com/dataproc/interactive" ) -_VSCODE_SESSION_URI_BASE = ( - "vscode://googlecloudtools.datacloud/dataproc/sessions" -) +_VSCODE_EXTENSION_ID = "googlecloudtools.datacloud" +_VSCODE_SESSION_URI_BASE = f"vscode://{_VSCODE_EXTENSION_ID}/dataproc/sessions" def _build_session_details_url( region: Optional[str], project_id: Optional[str], session_id: str ) -> str: - if environment.is_vscode(): + if environment.is_vscode() and environment.is_vscode_extension_installed( + _VSCODE_EXTENSION_ID + ): 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}" diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py index d8ef2d6..62c994c 100644 --- a/tests/unit/test_environment.py +++ b/tests/unit/test_environment.py @@ -14,6 +14,7 @@ import os import importlib +import subprocess import unittest from unittest import mock @@ -54,6 +55,43 @@ def test_is_vscode_false(self): os.environ.pop("VSCODE_PID", None) self.assertFalse(environment.is_vscode()) + @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") + def test_is_vscode_extension_installed_true(self, mock_run): + mock_run.return_value = subprocess.CompletedProcess( + args=["code", "--list-extensions"], + returncode=0, + stdout="ms-python.python\ngooglecloudtools.datacloud\n", + ) + self.assertTrue( + environment.is_vscode_extension_installed( + "googlecloudtools.datacloud" + ) + ) + + @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") + def test_is_vscode_extension_installed_false_not_listed(self, mock_run): + mock_run.return_value = subprocess.CompletedProcess( + args=["code", "--list-extensions"], + returncode=0, + stdout="ms-python.python\n", + ) + self.assertFalse( + environment.is_vscode_extension_installed( + "googlecloudtools.datacloud" + ) + ) + + @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") + def test_is_vscode_extension_installed_false_code_cli_missing( + self, mock_run + ): + mock_run.side_effect = FileNotFoundError() + self.assertFalse( + environment.is_vscode_extension_installed( + "googlecloudtools.datacloud" + ) + ) + def test_is_jupyter_true(self): os.environ["JPY_PARENT_PID"] = "67890" self.assertTrue(environment.is_jupyter()) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 02e8f9a..c74e8d4 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -1259,6 +1259,10 @@ 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( + "google.cloud.managed_spark_connect.session.environment.is_vscode_extension_installed", + return_value=True, + ) @mock.patch( "IPython.core.interactiveshell.InteractiveShell.initialized", return_value=True, @@ -1268,6 +1272,7 @@ def test_display_session_link_on_creation_vscode( self, mock_display, _mock_ipy, + _mock_ext_installed, ): mock.patch.dict( os.environ, @@ -1290,6 +1295,45 @@ def test_display_session_link_on_creation_vscode( ) self.assertNotIn(_MANAGED_SPARK_SESSIONS_BASE_URL, html_output) + @mock.patch( + "google.cloud.managed_spark_connect.session.environment.is_vscode_extension_installed", + return_value=False, + ) + @mock.patch( + "IPython.core.interactiveshell.InteractiveShell.initialized", + return_value=True, + ) + @mock.patch("IPython.display.display") + def test_display_session_link_on_creation_vscode_extension_not_installed( + self, + mock_display, + _mock_ipy, + _mock_ext_installed, + ): + 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( + f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/test-region/" + "test_session?project=test-project", + html_output, + ) + self.assertNotIn("vscode://", html_output) + + @mock.patch( + "google.cloud.managed_spark_connect.session.environment.is_vscode_extension_installed", + return_value=True, + ) @mock.patch.object(ManagedSparkSession, "getActiveSession") @mock.patch( "google.cloud.managed_spark_connect.session.get_active_s8s_session_response" @@ -1298,6 +1342,7 @@ def test_get_exiting_active_session_prints_vscode_url( self, mock_get_response, mock_get_active_session, + _mock_ext_installed, ): mock.patch.dict( os.environ, @@ -1325,7 +1370,13 @@ def test_get_exiting_active_session_prints_vscode_url( ) self.assertNotIn(_MANAGED_SPARK_SESSIONS_BASE_URL, printed) - def test_repr_html_uses_vscode_url_for_session_link(self): + @mock.patch( + "google.cloud.managed_spark_connect.session.environment.is_vscode_extension_installed", + return_value=True, + ) + def test_repr_html_uses_vscode_url_for_session_link( + self, _mock_ext_installed + ): mock.patch.dict( os.environ, { From eb358255f2ebd6472a10cb2a1976306502d82c66 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 10 Sep 2026 15:35:25 -0700 Subject: [PATCH 07/10] fix: fall back to scanning extensions dirs when code CLI is unavailable code may not be on PATH even when VS Code and the extension are installed (e.g. macOS before running "Shell Command: Install 'code' command in PATH", or a remote/SSH session using .vscode-server). Scan the on-disk extensions directories as a fallback so the vscode:// link still shows up in those cases. --- .../managed_spark_connect/environment.py | 36 +++++++++++- tests/unit/test_environment.py | 57 ++++++++++++++++++- 2 files changed, 87 insertions(+), 6 deletions(-) diff --git a/google/cloud/managed_spark_connect/environment.py b/google/cloud/managed_spark_connect/environment.py index bfb8135..2a0b8bd 100644 --- a/google/cloud/managed_spark_connect/environment.py +++ b/google/cloud/managed_spark_connect/environment.py @@ -48,13 +48,43 @@ def _installed_vscode_extensions() -> frozenset: ) +def _vscode_extensions_dirs() -> Tuple[str, ...]: + home = os.path.expanduser("~") + return ( + os.path.join(home, ".vscode", "extensions"), + os.path.join(home, ".vscode-server", "extensions"), + os.path.join(home, ".vscode-insiders", "extensions"), + os.path.join(home, ".vscode-server-insiders", "extensions"), + ) + + +@functools.lru_cache(maxsize=None) +def _is_extension_dir_present(extension_id: str) -> bool: + prefix = f"{extension_id}-" + for extensions_dir in _vscode_extensions_dirs(): + try: + entries = os.listdir(extensions_dir) + except OSError: + continue + if any(name.lower().startswith(prefix) for name in entries): + return True + return False + + def is_vscode_extension_installed(extension_id: str) -> bool: """True if the given VS Code extension id is installed. - Fails closed (returns False) if the `code` CLI is unavailable, times - out, or exits non-zero. + Checks via the `code` CLI first, falling back to scanning the + on-disk extensions directories directly, since `code` may not be on + PATH even when VS Code and the extension are installed (e.g. macOS + without "Shell Command: Install 'code' command in PATH" run, or a + remote/SSH session). Fails closed (returns False) if neither check + finds it. """ - return extension_id.lower() in _installed_vscode_extensions() + extension_id = extension_id.lower() + return extension_id in _installed_vscode_extensions() or ( + _is_extension_dir_present(extension_id) + ) def is_jupyter() -> bool: diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py index 62c994c..ee380a2 100644 --- a/tests/unit/test_environment.py +++ b/tests/unit/test_environment.py @@ -55,43 +55,94 @@ def test_is_vscode_false(self): os.environ.pop("VSCODE_PID", None) self.assertFalse(environment.is_vscode()) + @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") - def test_is_vscode_extension_installed_true(self, mock_run): + def test_is_vscode_extension_installed_true_via_cli( + self, mock_run, mock_listdir + ): mock_run.return_value = subprocess.CompletedProcess( args=["code", "--list-extensions"], returncode=0, stdout="ms-python.python\ngooglecloudtools.datacloud\n", ) + mock_listdir.side_effect = OSError() self.assertTrue( environment.is_vscode_extension_installed( "googlecloudtools.datacloud" ) ) + @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") - def test_is_vscode_extension_installed_false_not_listed(self, mock_run): + def test_is_vscode_extension_installed_false_not_listed( + self, mock_run, mock_listdir + ): mock_run.return_value = subprocess.CompletedProcess( args=["code", "--list-extensions"], returncode=0, stdout="ms-python.python\n", ) + mock_listdir.side_effect = OSError() self.assertFalse( environment.is_vscode_extension_installed( "googlecloudtools.datacloud" ) ) + @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") def test_is_vscode_extension_installed_false_code_cli_missing( - self, mock_run + self, mock_run, mock_listdir ): mock_run.side_effect = FileNotFoundError() + mock_listdir.side_effect = OSError() self.assertFalse( environment.is_vscode_extension_installed( "googlecloudtools.datacloud" ) ) + @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") + @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") + def test_is_vscode_extension_installed_true_via_disk_scan( + self, mock_run, mock_listdir + ): + mock_run.side_effect = FileNotFoundError() + + def fake_listdir(path): + if path.endswith(os.path.join(".vscode", "extensions")): + return [ + "ms-python.python-2024.1.0", + "googlecloudtools.datacloud-1.2.3", + ] + raise OSError() + + mock_listdir.side_effect = fake_listdir + self.assertTrue( + environment.is_vscode_extension_installed( + "googlecloudtools.datacloud" + ) + ) + + @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") + @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") + def test_is_vscode_extension_installed_true_via_remote_extensions_dir( + self, mock_run, mock_listdir + ): + mock_run.side_effect = FileNotFoundError() + + def fake_listdir(path): + if path.endswith(os.path.join(".vscode-server", "extensions")): + return ["googlecloudtools.datacloud-1.2.3"] + raise OSError() + + mock_listdir.side_effect = fake_listdir + self.assertTrue( + environment.is_vscode_extension_installed( + "googlecloudtools.datacloud" + ) + ) + def test_is_jupyter_true(self): os.environ["JPY_PARENT_PID"] = "67890" self.assertTrue(environment.is_jupyter()) From 6bb6ffa9e34ffab6c47444e5983b2dea05a9f72e Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 10 Sep 2026 16:08:48 -0700 Subject: [PATCH 08/10] fix: ignore the code CLI check on a vscode-server (Remote-SSH) backend On a remote backend, code runs through a client-forwarding shim whose --list-extensions behavior isn't reliable, so only the on-disk ~/.vscode-server/extensions scan is trusted there. --- .../managed_spark_connect/environment.py | 20 ++++++-- tests/unit/test_environment.py | 51 ++++++++++++++++--- 2 files changed, 61 insertions(+), 10 deletions(-) diff --git a/google/cloud/managed_spark_connect/environment.py b/google/cloud/managed_spark_connect/environment.py index 2a0b8bd..eb9bdf2 100644 --- a/google/cloud/managed_spark_connect/environment.py +++ b/google/cloud/managed_spark_connect/environment.py @@ -71,17 +71,31 @@ def _is_extension_dir_present(extension_id: str) -> bool: return False +def _is_remote_vscode_server() -> bool: + """True if this process is running as a vscode-server backend + (Remote-SSH, Tunnels, Dev Containers), as opposed to local desktop + VS Code. + """ + return os.path.isdir( + os.path.join(os.path.expanduser("~"), ".vscode-server") + ) + + def is_vscode_extension_installed(extension_id: str) -> bool: """True if the given VS Code extension id is installed. Checks via the `code` CLI first, falling back to scanning the on-disk extensions directories directly, since `code` may not be on PATH even when VS Code and the extension are installed (e.g. macOS - without "Shell Command: Install 'code' command in PATH" run, or a - remote/SSH session). Fails closed (returns False) if neither check - finds it. + without "Shell Command: Install 'code' command in PATH" run). On a + vscode-server (Remote-SSH) backend, `code` runs through a + client-forwarding shim with unreliable `--list-extensions` behavior, + so only the disk scan is used there. Fails closed (returns False) + if nothing finds it. """ extension_id = extension_id.lower() + if _is_remote_vscode_server(): + return _is_extension_dir_present(extension_id) return extension_id in _installed_vscode_extensions() or ( _is_extension_dir_present(extension_id) ) diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py index ee380a2..2a2ae2e 100644 --- a/tests/unit/test_environment.py +++ b/tests/unit/test_environment.py @@ -55,11 +55,13 @@ def test_is_vscode_false(self): os.environ.pop("VSCODE_PID", None) self.assertFalse(environment.is_vscode()) + @mock.patch("google.cloud.managed_spark_connect.environment.os.path.isdir") @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") def test_is_vscode_extension_installed_true_via_cli( - self, mock_run, mock_listdir + self, mock_run, mock_listdir, mock_isdir ): + mock_isdir.return_value = False mock_run.return_value = subprocess.CompletedProcess( args=["code", "--list-extensions"], returncode=0, @@ -72,11 +74,13 @@ def test_is_vscode_extension_installed_true_via_cli( ) ) + @mock.patch("google.cloud.managed_spark_connect.environment.os.path.isdir") @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") def test_is_vscode_extension_installed_false_not_listed( - self, mock_run, mock_listdir + self, mock_run, mock_listdir, mock_isdir ): + mock_isdir.return_value = False mock_run.return_value = subprocess.CompletedProcess( args=["code", "--list-extensions"], returncode=0, @@ -89,11 +93,13 @@ def test_is_vscode_extension_installed_false_not_listed( ) ) + @mock.patch("google.cloud.managed_spark_connect.environment.os.path.isdir") @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") def test_is_vscode_extension_installed_false_code_cli_missing( - self, mock_run, mock_listdir + self, mock_run, mock_listdir, mock_isdir ): + mock_isdir.return_value = False mock_run.side_effect = FileNotFoundError() mock_listdir.side_effect = OSError() self.assertFalse( @@ -102,11 +108,13 @@ def test_is_vscode_extension_installed_false_code_cli_missing( ) ) + @mock.patch("google.cloud.managed_spark_connect.environment.os.path.isdir") @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") def test_is_vscode_extension_installed_true_via_disk_scan( - self, mock_run, mock_listdir + self, mock_run, mock_listdir, mock_isdir ): + mock_isdir.return_value = False mock_run.side_effect = FileNotFoundError() def fake_listdir(path): @@ -124,12 +132,18 @@ def fake_listdir(path): ) ) + @mock.patch("google.cloud.managed_spark_connect.environment.os.path.isdir") @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") - def test_is_vscode_extension_installed_true_via_remote_extensions_dir( - self, mock_run, mock_listdir + def test_is_vscode_extension_installed_true_on_remote_server( + self, mock_run, mock_listdir, mock_isdir ): - mock_run.side_effect = FileNotFoundError() + mock_isdir.return_value = True + mock_run.return_value = subprocess.CompletedProcess( + args=["code", "--list-extensions"], + returncode=0, + stdout="", + ) def fake_listdir(path): if path.endswith(os.path.join(".vscode-server", "extensions")): @@ -143,6 +157,29 @@ def fake_listdir(path): ) ) + @mock.patch("google.cloud.managed_spark_connect.environment.os.path.isdir") + @mock.patch("google.cloud.managed_spark_connect.environment.os.listdir") + @mock.patch("google.cloud.managed_spark_connect.environment.subprocess.run") + def test_is_vscode_extension_installed_ignores_cli_on_remote_server( + self, mock_run, mock_listdir, mock_isdir + ): + # On a vscode-server backend, the `code` CLI runs through a + # client-forwarding shim with unreliable --list-extensions + # behavior, so its (misleading, here) positive result must be + # ignored in favor of the disk scan. + mock_isdir.return_value = True + mock_run.return_value = subprocess.CompletedProcess( + args=["code", "--list-extensions"], + returncode=0, + stdout="googlecloudtools.datacloud\n", + ) + mock_listdir.side_effect = OSError() + self.assertFalse( + environment.is_vscode_extension_installed( + "googlecloudtools.datacloud" + ) + ) + def test_is_jupyter_true(self): os.environ["JPY_PARENT_PID"] = "67890" self.assertTrue(environment.is_jupyter()) From 4eed13ddd65bcc8559c7c4860ee7dca79ee8a29d Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Sat, 12 Sep 2026 10:49:37 -0700 Subject: [PATCH 09/10] refactor: always show both console and vscode links instead of detecting Reliably detecting whether the Data Analytics Kit extension is installed proved impractical across local VS Code, Remote-SSH, and forks. Simplify by always showing both the Cloud Console link and the vscode:// deep link whenever running inside VS Code, and let the extension's own URI handler be the point of failure if it's missing. --- .../managed_spark_connect/environment.py | 74 ---------- google/cloud/managed_spark_connect/session.py | 50 ++++--- tests/unit/test_environment.py | 126 ------------------ tests/unit/test_session.py | 85 +++++------- 4 files changed, 63 insertions(+), 272 deletions(-) diff --git a/google/cloud/managed_spark_connect/environment.py b/google/cloud/managed_spark_connect/environment.py index eb9bdf2..e19dd97 100644 --- a/google/cloud/managed_spark_connect/environment.py +++ b/google/cloud/managed_spark_connect/environment.py @@ -12,9 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import functools import os -import subprocess import sys from typing import Callable, Tuple, List @@ -29,78 +27,6 @@ def is_vscode() -> bool: return os.getenv("VSCODE_PID") is not None -@functools.lru_cache(maxsize=1) -def _installed_vscode_extensions() -> frozenset: - try: - result = subprocess.run( - ["code", "--list-extensions"], - capture_output=True, - text=True, - timeout=5, - check=True, - ) - except (OSError, subprocess.SubprocessError): - return frozenset() - return frozenset( - line.strip().lower() - for line in result.stdout.splitlines() - if line.strip() - ) - - -def _vscode_extensions_dirs() -> Tuple[str, ...]: - home = os.path.expanduser("~") - return ( - os.path.join(home, ".vscode", "extensions"), - os.path.join(home, ".vscode-server", "extensions"), - os.path.join(home, ".vscode-insiders", "extensions"), - os.path.join(home, ".vscode-server-insiders", "extensions"), - ) - - -@functools.lru_cache(maxsize=None) -def _is_extension_dir_present(extension_id: str) -> bool: - prefix = f"{extension_id}-" - for extensions_dir in _vscode_extensions_dirs(): - try: - entries = os.listdir(extensions_dir) - except OSError: - continue - if any(name.lower().startswith(prefix) for name in entries): - return True - return False - - -def _is_remote_vscode_server() -> bool: - """True if this process is running as a vscode-server backend - (Remote-SSH, Tunnels, Dev Containers), as opposed to local desktop - VS Code. - """ - return os.path.isdir( - os.path.join(os.path.expanduser("~"), ".vscode-server") - ) - - -def is_vscode_extension_installed(extension_id: str) -> bool: - """True if the given VS Code extension id is installed. - - Checks via the `code` CLI first, falling back to scanning the - on-disk extensions directories directly, since `code` may not be on - PATH even when VS Code and the extension are installed (e.g. macOS - without "Shell Command: Install 'code' command in PATH" run). On a - vscode-server (Remote-SSH) backend, `code` runs through a - client-forwarding shim with unreliable `--list-extensions` behavior, - so only the disk scan is used there. Fails closed (returns False) - if nothing finds it. - """ - extension_id = extension_id.lower() - if _is_remote_vscode_server(): - return _is_extension_dir_present(extension_id) - return extension_id in _installed_vscode_extensions() or ( - _is_extension_dir_present(extension_id) - ) - - def is_jupyter() -> bool: """True if running in a Jupyter environment.""" return os.getenv("JPY_PARENT_PID") is not None diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index 258605a..807ec00 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -27,7 +27,7 @@ import tqdm from packaging import version from types import MethodType -from typing import Any, cast, ClassVar, Dict, Iterable, Optional, Union +from typing import Any, cast, ClassVar, Dict, Iterable, List, Optional, Tuple, Union from google.api_core import retry from google.api_core.client_options import ClientOptions @@ -71,18 +71,20 @@ "https://console.cloud.google.com/dataproc/interactive" ) -_VSCODE_EXTENSION_ID = "googlecloudtools.datacloud" -_VSCODE_SESSION_URI_BASE = f"vscode://{_VSCODE_EXTENSION_ID}/dataproc/sessions" +_VSCODE_SESSION_URI_BASE = ( + "vscode://googlecloudtools.datacloud/dataproc/sessions" +) -def _build_session_details_url( +def _build_session_details_links( region: Optional[str], project_id: Optional[str], session_id: str -) -> str: - if environment.is_vscode() and environment.is_vscode_extension_installed( - _VSCODE_EXTENSION_ID - ): - 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}" +) -> List[Tuple[str, str]]: + console_url = f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{region}/{session_id}?project={project_id}" + links = [("Managed Spark Session", console_url)] + if environment.is_vscode(): + vscode_url = f"{_VSCODE_SESSION_URI_BASE}/{session_id}?project={project_id}&location={region}" + links.append(("Managed Spark Session (Data Analytics Kit)", vscode_url)) + return links def _is_valid_label_value(value: str) -> bool: @@ -519,11 +521,11 @@ def _wait_for_session_available( ) def _display_session_link_on_creation(self, session_id): - session_url = _build_session_details_url( + links = _build_session_details_links( self._region, self._project_id, session_id ) - plain_message = ( - f"Creating Managed Spark Connect Session: {session_url}" + plain_message = "Creating Managed Spark Connect Session:\n" + ( + "\n".join(f" {label}: {url}" for label, url in links) ) if environment.is_colab_enterprise(): html_element = f""" @@ -532,10 +534,14 @@ def _display_session_link_on_creation(self, session_id): """ else: + links_html = "\n".join( + f'' + for label, url in links + ) html_element = f""" """ self._output_element_or_message(plain_message, html_element) @@ -588,11 +594,15 @@ def _get_exiting_active_session( session = ManagedSparkSession._default_session if session_response is not None: - session_url = _build_session_details_url( + links = _build_session_details_links( self._region, self._project_id, s8s_session_id ) + links_message = "\n".join( + f" {label}: {url}" for label, url in links + ) print( - f"Using existing Managed Spark Session (configuration changes may not be applied): {session_url}" + "Using existing Managed Spark Session (configuration " + f"changes may not be applied):\n{links_message}" ) self._display_view_session_details_button(s8s_session_id) if session is None: @@ -1125,9 +1135,13 @@ def _repr_html_(self) -> str: