From 50a72565db4377cbe6c6bd24e22949d06040a271 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Fri, 18 Sep 2026 14:29:51 -0700 Subject: [PATCH 1/4] feat: Automatically load interactive extras on import with opt-out controls Automatically initialize interactive extras when importing google.cloud.managed_spark_connect inside an IPython kernel: - Inject colabsqlviz's explore_dataframe() into IPython user_ns - Load the %dpip line magic extension (google.cloud.managed_spark_magics) - Package interactive dependencies (google-colabsqlviz>=0.3.0, ipython>=8.0, traitlets>=5.1) under the optional [interactive] extra (pip install 'google-cloud-spark-connect[interactive]') so non-interactive installs incur zero extra dependencies Conditional runtime initialization: - _ipython is only imported when traitlets is installed, and each extra is initialized only if its dependency is installed (importlib.util.find_spec). Uninstalled extras are skipped silently; only installed-but-broken extras produce a warning. - Because the active set of extras depends on what is installed in the environment, print a single summary line listing whichever extras were newly loaded: Loaded interactive extras: explore_dataframe(), %dpip. To disable: %config ManagedSparkConnect.enable_extras = False Extras that were already present are not listed. Add opt-out and runtime configuration controls: - Environment variable: MANAGED_SPARK_CONNECT_ENABLE_EXTRAS=false - IPython traitlet: ManagedSparkConnect.enable_extras = False (supports both persistent file-based config and runtime toggling via %config, tracking and only undoing changes that managed_spark_connect itself performed). --- DEVELOPING.md | 10 + README.md | 49 ++ .../cloud/managed_spark_connect/__init__.py | 12 + .../cloud/managed_spark_connect/_ipython.py | 421 ++++++++++ google/cloud/managed_spark_magics/__init__.py | 9 + google/cloud/managed_spark_magics/magics.py | 6 +- requirements-dev.txt | 2 + setup.py | 12 + tests/unit/test_init.py | 791 ++++++++++++++++++ 9 files changed, 1311 insertions(+), 1 deletion(-) create mode 100644 google/cloud/managed_spark_connect/_ipython.py diff --git a/DEVELOPING.md b/DEVELOPING.md index c9a8a5ea..b4a6df2b 100644 --- a/DEVELOPING.md +++ b/DEVELOPING.md @@ -61,6 +61,16 @@ pytest Tests that require magic functionality will be automatically skipped if the dependencies are not available. +## Testing the Interactive Extras + +The dependencies of the interactive extras (`explore_dataframe`, `%dpip`) live +in the `interactive` extra (`pip install '.[interactive]'`), and +`requirements-dev.txt` already includes them. The tests in +`tests/unit/test_init.py` assume `google-colabsqlviz`, `ipython` and +`traitlets` are importable and will fail, not skip, if they are not. Missing +dependencies are simulated within the tests themselves, by setting +`sys.modules[name] = None`. + The integration tests in particular can take a while to run. To speed up the testing cycle, you can run them in parallel. You can do so using the `xdist` plugin by setting the `-n` flag to the number of parallel runners you want to diff --git a/README.md b/README.md index c1117926..a9cdf3f3 100644 --- a/README.md +++ b/README.md @@ -121,6 +121,55 @@ To create or connect to a named session: 5. A session with a given ID that is in a TERMINATED state cannot be reused. It must be deleted before a new session with the same ID can be created. +### Interactive Extras + +For notebooks and other interactive use, install the `interactive` extra: + +```sh +pip install 'google-cloud-spark-connect[interactive]' +``` + +When you import the package inside an IPython kernel, it then automatically sets +up a few interactive conveniences: + +- `explore_dataframe()` from + [google-colabsqlviz](https://pypi.org/project/google-colabsqlviz/) is injected + into your notebook globals. +- The `%dpip` line magic is loaded. + +```python +import google.cloud.managed_spark_connect # extras load here +``` + +Each extra is only set up if its dependency is importable, so without the +`interactive` extra you get whichever of them your environment already happens +to provide. On import, a single line lists the extras that were loaded. + +The extras won't override anything you've already set up, for example if `explore_dataframe` is already present, it is left alone. + +#### Opting out + +Set the environment variable before starting the kernel: + +```sh +export MANAGED_SPARK_CONNECT_ENABLE_EXTRAS=false +``` + +Alternatively, configure it through IPython, either persistently in +`~/.ipython/profile_default/ipython_config.py`: + +```python +c.ManagedSparkConnect.enable_extras = False +``` + +Or at runtime: + +```python +%config ManagedSparkConnect.enable_extras = False +``` + +An explicit IPython setting takes precedence over the environment variable. + ### Using Spark SQL Magic Commands (Jupyter Notebooks) The package supports the [sparksql-magic](https://github.com/cryeo/sparksql-magic) library for executing Spark SQL queries directly in Jupyter notebooks. diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index 23567583..06082a24 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import importlib.metadata +import importlib.util import warnings from .session import ManagedSparkSession @@ -28,3 +29,14 @@ ) except Exception: pass + +# traitlets, like every other dependency of the interactive extras, only comes +# with the [interactive] extra. IPython depends on it, so without it there can +# be no shell and nothing to initialize. +if importlib.util.find_spec("traitlets") is not None: + # Only _init_extras: the ManagedSparkConnect configurable is looked up by + # class name (`c.ManagedSparkConnect...`), so it need not be importable + # from here, and would otherwise exist only when traitlets is installed. + from ._ipython import _init_extras + + _init_extras() diff --git a/google/cloud/managed_spark_connect/_ipython.py b/google/cloud/managed_spark_connect/_ipython.py new file mode 100644 index 00000000..e43f7079 --- /dev/null +++ b/google/cloud/managed_spark_connect/_ipython.py @@ -0,0 +1,421 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +import importlib.util +import logging +import os +import sys +from typing import Any, Dict, List, Optional, Set + +from traitlets import Bool, default, observe +from traitlets.config import Config, Configurable + +logger = logging.getLogger(__name__) + +ENV_ENABLE_EXTRAS = "MANAGED_SPARK_CONNECT_ENABLE_EXTRAS" + +_ENV_FALSE_VALUES = frozenset(("0", "false", "no", "off")) + +# Module that must be installed for explore_dataframe() to be injected. +_COLABSQLVIZ_MODULE = "google.colabsqlviz" + +# (extension module, magic kind, magic name). Each extension is only loaded if +# its module is installed. +_EXTRAS_EXTENSIONS = (("google.cloud.managed_spark_magics", "line", "dpip"),) + + +def _is_module_available(name: str) -> bool: + """Whether `name` is installed, without importing it. + + All of the extras' dependencies are optional (they come from the + [interactive] extra), so a module that isn't installed means that extra is + deliberately absent, not broken. Anything other than a clean "not found" is + reported as available so that the subsequent import surfaces the real error. + """ + try: + return importlib.util.find_spec(name) is not None + except ModuleNotFoundError: + # A parent package is missing. + return False + except Exception: + return True + + +def _is_env_extras_enabled() -> bool: + val = os.getenv(ENV_ENABLE_EXTRAS) + if val is None or not val.strip(): + # Treat an empty/whitespace-only value the same as an unset variable, + # so that `MANAGED_SPARK_CONNECT_ENABLE_EXTRAS=` is not silently + # interpreted as "enabled". + return True + return val.strip().lower() not in _ENV_FALSE_VALUES + + +class ManagedSparkConnect(Configurable): + """IPython configuration for google.cloud.managed_spark_connect extras.""" + + # No static default_value: the dynamic `_default_enable_extras` below + # always wins. + enable_extras = Bool( + help=( + "Whether to automatically load the installed interactive extras " + "(explore_dataframe, %dpip)." + ), + ).tag(config=True) + + def __init__(self, shell: Any = None, **kwargs: Any) -> None: + self._initialized = False + self._shell = shell + super().__init__(**kwargs) + self._initialized = True + + @default("enable_extras") + def _default_enable_extras(self) -> bool: + return _is_env_extras_enabled() + + @observe("enable_extras") + def _observe_enable_extras(self, change: Dict[str, Any]) -> None: + # Traitlets applies file-based config inside `super().__init__()`, and + # the resulting notification fires before the shell has been wired up. + # Skip it: `_init_extras` loads the extras once construction is done. + if not getattr(self, "_initialized", False): + return + shell = self._shell or self.parent + if shell is None: + return + if change["new"]: + _load_extras_for_shell(shell) + else: + _unload_extras_for_shell(shell) + + +@dataclass +class _ShellExtrasState: + shell: Any + injected_explore_dataframe: Any = None + loaded_extensions: Set[str] = field(default_factory=set) + config_instance: Optional[ManagedSparkConnect] = None + + +# Keyed by id() rather than being a WeakKeyDictionary. A weak key would not +# release anything here: the value reaches the key via +# `config_instance._shell` and traitlets' own `Configurable.parent`, and +# `InteractiveShell.__init__` hands a bound method to `atexit`, so a shell is +# retained for the life of the process regardless. `state.shell` is what makes +# id() keying safe, since CPython reuses addresses; see `_get_shell_state`. +_SHELL_STATES: Dict[int, _ShellExtrasState] = {} + + +def _get_shell_state(ip: Any) -> _ShellExtrasState: + shell_id = id(ip) + state = _SHELL_STATES.get(shell_id) + if state is None or state.shell is not ip: + state = _ShellExtrasState(shell=ip) + _SHELL_STATES[shell_id] = state + return state + + +def _is_same_class(obj: Any, cls: type) -> bool: + """Compares by module+qualname so that reloaded classes still match.""" + obj_cls = type(obj) + return ( + obj_cls.__module__ == cls.__module__ + and obj_cls.__qualname__ == cls.__qualname__ + ) + + +def _get_or_create_config(ip: Any) -> ManagedSparkConnect: + state = _get_shell_state(ip) + ip_config = getattr(ip, "config", None) + + if state.config_instance is None: + parent = ip if isinstance(ip, Configurable) else None + config = ( + ip_config + if (parent is None and isinstance(ip_config, Config)) + else None + ) + # Precedence is left entirely to traitlets: an explicit + # `c.ManagedSparkConnect.enable_extras` from the shell's config wins, + # otherwise `_default_enable_extras` consults the environment. + cfg = ManagedSparkConnect(shell=ip, parent=parent, config=config) + state.config_instance = cfg + else: + cfg = state.config_instance + cfg._shell = ip + + configurables = getattr(ip, "configurables", None) + if isinstance(configurables, list): + # Remove any stale ManagedSparkConnect instances (e.g. across module + # reloads) + configurables[:] = [ + c + for c in configurables + if c is cfg or not _is_same_class(c, ManagedSparkConnect) + ] + if cfg not in configurables: + configurables.append(cfg) + + return cfg + + +def _import_explore_dataframe() -> Any: + """Returns colabsqlviz's explore_dataframe(), or None if unavailable.""" + try: + from google.colabsqlviz.explore_dataframe import explore_dataframe + + return explore_dataframe + except Exception: + logger.debug("Failed to import explore_dataframe", exc_info=True) + return None + + +def _safe_repr(value: Any, max_len: int = 60) -> str: + """repr() that tolerates objects whose __repr__ raises.""" + try: + val_repr = repr(value) + except Exception: + return f"<{type(value).__name__} with a failing __repr__>" + if len(val_repr) > max_len: + return f"{val_repr[:max_len]}..." + return val_repr + + +def _print_extras_failure() -> None: + """Tells the user something went wrong, without diagnosing what. + + Every individual failure below is logged at debug level and otherwise + swallowed, which leaves a broken install looking exactly like extras being + switched off. One line, the same regardless of cause, is enough to point + somebody at the logs. + """ + print( + "\033[94m⚠️ [google.cloud.managed_spark_connect]\033[0m" + " Failed to load interactive extras." + " For details, enable debug logging and restart the kernel." + ) + + +def _is_extension_loaded( + ip: Any, ext_name: str, magic_kind: str, magic_name: str +) -> bool: + mod = sys.modules.get(ext_name) + if mod is not None and getattr( + getattr(mod, "__spec__", None), "_initializing", False + ): + return True + + ext_mgr = getattr(ip, "extension_manager", None) + loaded = getattr(ext_mgr, "loaded", None) + if isinstance(loaded, (set, dict, list, tuple)) and ext_name in loaded: + return True + + magics_mgr = getattr(ip, "magics_manager", None) + magics = getattr(magics_mgr, "magics", None) + if isinstance(magics, dict): + kind_dict = magics.get(magic_kind) + if isinstance(kind_dict, dict) and magic_name in kind_dict: + return True + + return False + + +def _print_extras_loaded(loaded: List[str]) -> None: + """Lists the extras that were just loaded, in a single line. + + Which extras get loaded depends on what happens to be installed, which the + user may not know (or may have forgotten), so name them explicitly rather + than leaving them to be discovered. + """ + names = ", ".join(f"\033[1m{name}\033[0m" for name in loaded) + print( + "\033[94m👉 [google.cloud.managed_spark_connect]\033[0m" + f" Loaded interactive extras: {names}." + " To disable: %config ManagedSparkConnect.enable_extras = False" + ) + + +def _magic_display_name(magic_kind: str, magic_name: str) -> str: + prefix = "%%" if magic_kind == "cell" else "%" + return f"{prefix}{magic_name}" + + +def _load_extras_for_shell(ip: Any) -> None: + state = _get_shell_state(ip) + target_name = "explore_dataframe" + failed = False + # Only what this call loaded; anything that was already present, whether + # from us or the user, is not news. + loaded: List[str] = [] + + # 1. Inject explore_dataframe from google-colabsqlviz + try: + user_ns = getattr(ip, "user_ns", None) + if not _is_module_available(_COLABSQLVIZ_MODULE): + logger.debug( + "Not injecting %s: %s is not installed", + target_name, + _COLABSQLVIZ_MODULE, + ) + elif isinstance(user_ns, dict): + if ( + state.injected_explore_dataframe is not None + and user_ns.get(target_name) is state.injected_explore_dataframe + ): + pass + elif target_name not in user_ns: + from google.colabsqlviz.explore_dataframe import ( + explore_dataframe, + ) + + ip.push({target_name: explore_dataframe}) + state.injected_explore_dataframe = explore_dataframe + loaded.append(f"{target_name}()") + else: + current_val = user_ns[target_name] + ours = _import_explore_dataframe() + if ours is not None and current_val is ours: + # Already bound to the exact function we would have + # injected, e.g. because the user imported it themselves. + # Warning about that would be pure noise. + pass + else: + val_repr = _safe_repr(current_val) + print( + "\033[94m⚠️ [google.cloud.managed_spark_connect]\033[0m" + f" Did not inject \033[1m{target_name}()\033[0m" + f" because it is already defined as: {val_repr}" + ) + except Exception: + logger.debug( + "Failed to inject %s into IPython user_ns", + target_name, + exc_info=True, + ) + failed = True + + # 2. Load magic extensions + ext_mgr = getattr(ip, "extension_manager", None) + if ext_mgr is not None and hasattr(ext_mgr, "load_extension"): + for ext_name, magic_kind, magic_name in _EXTRAS_EXTENSIONS: + try: + if ext_name in state.loaded_extensions or ( + _is_extension_loaded(ip, ext_name, magic_kind, magic_name) + ): + continue + if not _is_module_available(ext_name): + logger.debug( + "Not loading IPython extension %s: not installed", + ext_name, + ) + continue + res = ext_mgr.load_extension(ext_name) + if res is None: + state.loaded_extensions.add(ext_name) + loaded.append(_magic_display_name(magic_kind, magic_name)) + except Exception: + logger.debug( + "Failed to load IPython extension %s", + ext_name, + exc_info=True, + ) + failed = True + + if loaded: + _print_extras_loaded(loaded) + if failed: + _print_extras_failure() + + +def _unload_extras_for_shell(ip: Any) -> None: + state = _get_shell_state(ip) + target_name = "explore_dataframe" + + # 1. Remove explore_dataframe only if we injected it and it wasn't overwritten + try: + if state.injected_explore_dataframe is not None: + user_ns = getattr(ip, "user_ns", None) + if ( + isinstance(user_ns, dict) + and user_ns.get(target_name) is state.injected_explore_dataframe + ): + del user_ns[target_name] + state.injected_explore_dataframe = None + except Exception: + logger.debug( + "Failed to remove %s from IPython user_ns", + target_name, + exc_info=True, + ) + + # 2. Unload only the extensions that we loaded + ext_mgr = getattr(ip, "extension_manager", None) + magics_mgr = getattr(ip, "magics_manager", None) + magics = getattr(magics_mgr, "magics", None) + loaded = getattr(ext_mgr, "loaded", None) + + for ext_name, magic_kind, magic_name in _EXTRAS_EXTENSIONS: + if ext_name in state.loaded_extensions: + try: + if ext_mgr is not None and hasattr(ext_mgr, "unload_extension"): + ext_mgr.unload_extension(ext_name) + except Exception: + logger.debug( + "Failed to unload IPython extension %s", + ext_name, + exc_info=True, + ) + try: + if isinstance(magics, dict): + kind_dict = magics.get(magic_kind) + if isinstance(kind_dict, dict): + kind_dict.pop(magic_name, None) + if isinstance(loaded, set): + loaded.discard(ext_name) + except Exception: + logger.debug( + "Failed to clean up magic %s for extension %s", + magic_name, + ext_name, + exc_info=True, + ) + state.loaded_extensions.discard(ext_name) + + +def _init_extras() -> None: + """Initialize interactive extras on package import unless opted out.""" + ip = None + try: + from IPython import get_ipython + + ip = get_ipython() + if ip is None: + return + + # `cfg.enable_extras` already folds in both the environment variable + # (via the trait's dynamic default) and any explicit configuration, and + # it is the same value `%config` toggles later, so it is the only thing + # that should be consulted here. + cfg = _get_or_create_config(ip) + if cfg.enable_extras: + _load_extras_for_shell(ip) + except Exception: + logger.debug( + "Failed to initialize IPython extras on import", exc_info=True + ) + if ip is not None: + # Only worth saying inside a shell: outside one there are no extras + # to miss, and printing would be noise in ordinary scripts. + _print_extras_failure() diff --git a/google/cloud/managed_spark_magics/__init__.py b/google/cloud/managed_spark_magics/__init__.py index 79632f57..eb667258 100644 --- a/google/cloud/managed_spark_magics/__init__.py +++ b/google/cloud/managed_spark_magics/__init__.py @@ -17,3 +17,12 @@ def load_ipython_extension(ipython): ipython.register_magics(ManagedSparkMagics) + + +def unload_ipython_extension(ipython): + magics_mgr = getattr(ipython, "magics_manager", None) + magics = getattr(magics_mgr, "magics", None) + if isinstance(magics, dict): + line_magics = magics.get("line") + if isinstance(line_magics, dict): + line_magics.pop("dpip", None) diff --git a/google/cloud/managed_spark_magics/magics.py b/google/cloud/managed_spark_magics/magics.py index 54363ae1..38ca3c7c 100644 --- a/google/cloud/managed_spark_magics/magics.py +++ b/google/cloud/managed_spark_magics/magics.py @@ -16,7 +16,6 @@ import shlex from IPython.core.magic import (Magics, magics_class, line_magic) -from google.cloud.managed_spark_connect import ManagedSparkSession @magics_class @@ -35,6 +34,11 @@ def dpip(self, line): Custom magic to install pip packages as Spark Connect artifacts. Usage: %dpip install pandas numpy """ + # Imported here rather than at module scope: google.cloud + # .managed_spark_connect loads this extension from its own __init__, + # so a module-level import back into it would be circular. + from google.cloud.managed_spark_connect import ManagedSparkSession + try: args = shlex.split(line) diff --git a/requirements-dev.txt b/requirements-dev.txt index 5cf7026e..09806396 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,5 +1,6 @@ google-api-core>=2.19 google-cloud-dataproc>=5.18 +google-colabsqlviz>=0.3.0 ipython~=9.1 ipywidgets>=8.0.0 packaging>=20.0 @@ -8,4 +9,5 @@ pyspark[connect]~=4.0.0 setuptools>=72.0 sparksql-magic>=0.0.3 tqdm>=4.67 +traitlets>=5.1 websockets>=14.0 diff --git a/setup.py b/setup.py index 89637d54..b1850daa 100644 --- a/setup.py +++ b/setup.py @@ -36,4 +36,16 @@ "tqdm>=4.67", "websockets>=14.0", ], + extras_require={ + # Everything here is optional at runtime: managed_spark_connect._ipython + # only initializes the extras whose dependencies are importable. + "interactive": [ + "google-colabsqlviz>=0.3.0", + # Imported directly by managed_spark_connect._ipython and + # managed_spark_magics; previously these only arrived transitively + # via google-colabsqlviz. + "ipython>=8.0", + "traitlets>=5.1", + ], + }, ) diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index 794f6fa8..19422110 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -135,5 +135,796 @@ def test_invalid_runtime_version_logs_warning(self, mock_logger): ) +class TestDependencies(unittest.TestCase): + + def test_colabsqlviz_available(self): + import importlib.metadata + from packaging.version import Version + import google.colabsqlviz + from google.colabsqlviz.explore_dataframe import ( + display, + explore_dataframe, + ) + + self.assertIsNotNone(google.colabsqlviz) + self.assertTrue(callable(display)) + self.assertTrue(callable(explore_dataframe)) + version = importlib.metadata.version("google-colabsqlviz") + self.assertGreaterEqual(Version(version), Version("0.3.0")) + + def test_managed_spark_magics_available(self): + import google.cloud.managed_spark_magics as magics_ext + + self.assertTrue(callable(magics_ext.load_ipython_extension)) + self.assertTrue(callable(magics_ext.unload_ipython_extension)) + + +def _make_mock_shell(user_ns=None, config=None): + """Builds a mock shell whose push() mutates user_ns, like the real one. + + Without the side effect, mock shells silently swallow ip.push() and only + pass if the code under test also writes to user_ns directly. + """ + mock_ip = mock.Mock() + mock_ip.user_ns = {} if user_ns is None else user_ns + mock_ip.config = config + mock_ip.configurables = [] + mock_ip.push.side_effect = mock_ip.user_ns.update + mock_ip.extension_manager.load_extension.return_value = None + return mock_ip + + +class TestExtrasUnit(unittest.TestCase): + + def setUp(self): + import os + from google.cloud.managed_spark_connect import _ipython + + env_patcher = mock.patch.dict(os.environ) + env_patcher.start() + self.addCleanup(env_patcher.stop) + os.environ.pop(_ipython.ENV_ENABLE_EXTRAS, None) + _ipython._SHELL_STATES.clear() + self.addCleanup(_ipython._SHELL_STATES.clear) + + @mock.patch("sys.stdout", new_callable=unittest.mock.MagicMock) + @mock.patch("IPython.get_ipython", return_value=None) + def test_no_extras_when_not_in_ipython(self, mock_get_ipython, mock_stdout): + from google.cloud.managed_spark_connect._ipython import _init_extras + + _init_extras() + mock_get_ipython.assert_called_once() + mock_stdout.write.assert_not_called() + + @mock.patch("IPython.get_ipython") + def test_successful_extras_loading(self, mock_get_ipython): + import io + from google.colabsqlviz.explore_dataframe import explore_dataframe + from google.cloud.managed_spark_connect._ipython import _init_extras + + mock_ip = _make_mock_shell() + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + mock_ip.push.assert_called_once_with( + {"explore_dataframe": explore_dataframe} + ) + self.assertIs(mock_ip.user_ns["explore_dataframe"], explore_dataframe) + mock_ip.extension_manager.load_extension.assert_called_once_with( + "google.cloud.managed_spark_magics" + ) + output = stdout_capture.getvalue() + self.assertIn("[google.cloud.managed_spark_connect]", output) + self.assertIn("Loaded interactive extras", output) + self.assertIn("explore_dataframe()", output) + self.assertIn("%dpip", output) + self.assertIn("ManagedSparkConnect.enable_extras = False", output) + # One summary line, not one per extra. + self.assertEqual(len(output.splitlines()), 1) + + @mock.patch("IPython.get_ipython") + def test_fallback_when_already_defined_short_repr(self, mock_get_ipython): + import io + from google.cloud.managed_spark_connect._ipython import ( + _get_or_create_config, + _init_extras, + ) + + mock_ip = _make_mock_shell(user_ns={"explore_dataframe": 42}) + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + mock_ip.push.assert_not_called() + output = stdout_capture.getvalue() + self.assertIn("[google.cloud.managed_spark_connect]", output) + self.assertIn("Did not inject", output) + self.assertIn("explore_dataframe()", output) + self.assertIn("because it is already defined as: 42", output) + self.assertNotIn("42...", output) + + # Disabling extras via traitlet must not remove user-defined explore_dataframe + _get_or_create_config(mock_ip).enable_extras = False + self.assertEqual(mock_ip.user_ns.get("explore_dataframe"), 42) + + @mock.patch("IPython.get_ipython") + def test_fallback_when_already_defined_long_repr(self, mock_get_ipython): + import io + from google.cloud.managed_spark_connect._ipython import _init_extras + + long_val = "x" * 100 + mock_ip = _make_mock_shell(user_ns={"explore_dataframe": long_val}) + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + mock_ip.push.assert_not_called() + output = stdout_capture.getvalue() + expected_repr = f"{repr(long_val)[:60]}..." + self.assertIn("[google.cloud.managed_spark_connect]", output) + self.assertIn("Did not inject", output) + self.assertIn("explore_dataframe()", output) + self.assertIn( + f"because it is already defined as: {expected_repr}", + output, + ) + + @mock.patch("IPython.get_ipython") + def test_fallback_when_already_defined_unreprable(self, mock_get_ipython): + import io + from google.cloud.managed_spark_connect._ipython import _init_extras + + class Unreprable: + + def __repr__(self): + raise ValueError("no repr for you") + + mock_ip = _make_mock_shell(user_ns={"explore_dataframe": Unreprable()}) + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + mock_ip.push.assert_not_called() + output = stdout_capture.getvalue() + self.assertIn("Did not inject", output) + self.assertIn("Unreprable", output) + # A broken __repr__ must not abort the rest of the initialization. + mock_ip.extension_manager.load_extension.assert_called_once_with( + "google.cloud.managed_spark_magics" + ) + + @mock.patch("IPython.get_ipython") + def test_no_warning_when_already_defined_same_function( + self, mock_get_ipython + ): + import io + from google.colabsqlviz.explore_dataframe import explore_dataframe + from google.cloud.managed_spark_connect._ipython import _init_extras + + mock_ip = _make_mock_shell( + user_ns={"explore_dataframe": explore_dataframe} + ) + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + # The name already refers to the exact function we would have injected, + # so there is nothing to warn the user about, and since we didn't load + # it, it isn't listed either. + mock_ip.push.assert_not_called() + output = stdout_capture.getvalue() + self.assertNotIn("Did not inject", output) + self.assertNotIn("explore_dataframe", output) + self.assertIn("Loaded interactive extras", output) + self.assertIs(mock_ip.user_ns["explore_dataframe"], explore_dataframe) + + @mock.patch("IPython.get_ipython") + def test_opt_out_via_env_var(self, mock_get_ipython): + import os + from google.cloud.managed_spark_connect._ipython import ( + ENV_ENABLE_EXTRAS, + _init_extras, + ) + + for falsy_val in ("false", "FALSE", "0", "no", "off", " false "): + with self.subTest(val=falsy_val): + mock_ip = _make_mock_shell() + mock_get_ipython.return_value = mock_ip + + with mock.patch.dict( + os.environ, {ENV_ENABLE_EXTRAS: falsy_val} + ): + _init_extras() + + mock_ip.push.assert_not_called() + mock_ip.extension_manager.load_extension.assert_not_called() + self.assertNotIn("explore_dataframe", mock_ip.user_ns) + + @mock.patch("IPython.get_ipython") + def test_opt_out_via_traitlet_config(self, mock_get_ipython): + from traitlets.config import Config + from google.cloud.managed_spark_connect._ipython import _init_extras + + cfg = Config() + cfg.ManagedSparkConnect.enable_extras = False + mock_ip = _make_mock_shell(config=cfg) + mock_get_ipython.return_value = mock_ip + + _init_extras() + + mock_ip.push.assert_not_called() + mock_ip.extension_manager.load_extension.assert_not_called() + self.assertNotIn("explore_dataframe", mock_ip.user_ns) + + @mock.patch("IPython.get_ipython") + def test_explicit_config_beats_env_var(self, mock_get_ipython): + import io + import os + from traitlets.config import Config + from google.cloud.managed_spark_connect._ipython import ( + ENV_ENABLE_EXTRAS, + _init_extras, + ) + + # An explicit `c.ManagedSparkConnect.enable_extras = True` is a more + # specific signal than the environment variable, so it must win. + cfg = Config() + cfg.ManagedSparkConnect.enable_extras = True + mock_ip = _make_mock_shell(config=cfg) + mock_get_ipython.return_value = mock_ip + + with mock.patch.dict(os.environ, {ENV_ENABLE_EXTRAS: "false"}): + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + + self.assertIn("explore_dataframe", mock_ip.user_ns) + mock_ip.extension_manager.load_extension.assert_called_once_with( + "google.cloud.managed_spark_magics" + ) + + @mock.patch("IPython.get_ipython") + def test_blank_env_var_is_treated_as_unset(self, mock_get_ipython): + import io + import os + from google.cloud.managed_spark_connect._ipython import ( + ENV_ENABLE_EXTRAS, + _init_extras, + ) + + for blank_val in ("", " "): + with self.subTest(val=blank_val): + mock_ip = _make_mock_shell() + mock_get_ipython.return_value = mock_ip + + with mock.patch.dict( + os.environ, {ENV_ENABLE_EXTRAS: blank_val} + ): + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + + self.assertIn("explore_dataframe", mock_ip.user_ns) + + @mock.patch("IPython.get_ipython") + def test_traitlet_disable_preserves_overwritten_explore_dataframe( + self, mock_get_ipython + ): + import io + from google.cloud.managed_spark_connect._ipython import ( + _get_or_create_config, + _init_extras, + ) + + mock_ip = _make_mock_shell() + mock_get_ipython.return_value = mock_ip + + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + + custom_fn = lambda df: "custom" + mock_ip.user_ns["explore_dataframe"] = custom_fn + + _get_or_create_config(mock_ip).enable_extras = False + self.assertIs(mock_ip.user_ns.get("explore_dataframe"), custom_fn) + mock_ip.extension_manager.unload_extension.assert_called_once_with( + "google.cloud.managed_spark_magics" + ) + + @mock.patch("IPython.get_ipython", side_effect=RuntimeError("Kernel error")) + def test_exception_silently_caught(self, mock_get_ipython): + from google.cloud.managed_spark_connect._ipython import _init_extras + + try: + _init_extras() + except Exception as e: + self.fail(f"_init_extras raised an exception: {e}") + + @mock.patch("IPython.get_ipython", side_effect=RuntimeError("Kernel error")) + def test_no_failure_message_outside_a_shell(self, mock_get_ipython): + import io + from google.cloud.managed_spark_connect._ipython import _init_extras + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + # There is no shell, so there are no extras to miss. + self.assertEqual(stdout_capture.getvalue(), "") + + @mock.patch("IPython.get_ipython") + def test_failure_message_printed_once(self, mock_get_ipython): + import io + from google.cloud.managed_spark_connect._ipython import _init_extras + + mock_ip = _make_mock_shell() + mock_ip.extension_manager.load_extension.side_effect = RuntimeError( + "boom" + ) + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + with mock.patch( + "google.cloud.managed_spark_connect._ipython" + "._import_explore_dataframe", + side_effect=RuntimeError("boom"), + ): + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + output = stdout_capture.getvalue() + self.assertIn("Failed to load interactive extras", output) + # Two separate failures (the injection and the extension) must not + # produce two separate messages. + self.assertEqual(output.count("Failed to load interactive extras"), 1) + + @mock.patch("IPython.get_ipython") + def test_failure_message_when_colabsqlviz_fails_to_import( + self, mock_get_ipython + ): + import builtins + import io + from google.cloud.managed_spark_connect._ipython import _init_extras + + mock_ip = _make_mock_shell() + mock_get_ipython.return_value = mock_ip + + real_import = builtins.__import__ + + def fail_colabsqlviz(name, *args, **kwargs): + if name.startswith("google.colabsqlviz"): + raise ImportError("no colabsqlviz here") + return real_import(name, *args, **kwargs) + + stdout_capture = io.StringIO() + with mock.patch("builtins.__import__", side_effect=fail_colabsqlviz): + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + # colabsqlviz is installed but broken, which is worth reporting. + self.assertNotIn("explore_dataframe", mock_ip.user_ns) + self.assertIn( + "Failed to load interactive extras", stdout_capture.getvalue() + ) + # The magics are independent of colabsqlviz and should still load. + mock_ip.extension_manager.load_extension.assert_called_once_with( + "google.cloud.managed_spark_magics" + ) + + @mock.patch("IPython.get_ipython") + def test_colabsqlviz_not_installed_is_skipped(self, mock_get_ipython): + import io + import sys + from google.cloud.managed_spark_connect._ipython import _init_extras + + mock_ip = _make_mock_shell() + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + # A None entry in sys.modules makes both find_spec() and import behave + # as if the module were not installed. + with mock.patch.dict(sys.modules, {"google.colabsqlviz": None}): + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + mock_ip.push.assert_not_called() + self.assertNotIn("explore_dataframe", mock_ip.user_ns) + mock_ip.extension_manager.load_extension.assert_called_once_with( + "google.cloud.managed_spark_magics" + ) + # Only what was actually loaded is listed, and a missing optional + # dependency is not a failure. + output = stdout_capture.getvalue() + self.assertEqual(len(output.splitlines()), 1) + self.assertIn("Loaded interactive extras", output) + self.assertIn("%dpip", output) + self.assertNotIn("explore_dataframe", output) + self.assertNotIn("Failed to load interactive extras", output) + + @mock.patch("IPython.get_ipython") + def test_colabsqlviz_not_installed_does_not_warn_about_existing_name( + self, mock_get_ipython + ): + import io + import sys + from google.cloud.managed_spark_connect._ipython import _init_extras + + mock_ip = _make_mock_shell(user_ns={"explore_dataframe": 42}) + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + with mock.patch.dict(sys.modules, {"google.colabsqlviz": None}): + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + # We would not have injected anything, so there is no conflict to + # warn about. + self.assertNotIn("Did not inject", stdout_capture.getvalue()) + self.assertEqual(mock_ip.user_ns["explore_dataframe"], 42) + + @mock.patch("IPython.get_ipython") + def test_nothing_printed_when_nothing_is_loaded(self, mock_get_ipython): + import io + import sys + from google.cloud.managed_spark_connect._ipython import _init_extras + + mock_ip = _make_mock_shell() + mock_get_ipython.return_value = mock_ip + + stdout_capture = io.StringIO() + with mock.patch.dict( + sys.modules, + { + "google.colabsqlviz": None, + "google.cloud.managed_spark_magics": None, + }, + ): + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + mock_ip.extension_manager.load_extension.assert_not_called() + self.assertEqual(stdout_capture.getvalue(), "") + + @mock.patch("IPython.get_ipython") + def test_summary_not_repeated_on_second_load(self, mock_get_ipython): + import io + from google.cloud.managed_spark_connect._ipython import ( + _init_extras, + _load_extras_for_shell, + ) + + mock_ip = _make_mock_shell() + mock_get_ipython.return_value = mock_ip + + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _load_extras_for_shell(mock_ip) + + # Everything is already loaded, so there is nothing new to report. + self.assertEqual(stdout_capture.getvalue(), "") + + def test_is_module_available(self): + from google.cloud.managed_spark_connect._ipython import ( + _is_module_available, + ) + + self.assertTrue(_is_module_available("json")) + self.assertTrue(_is_module_available("google.colabsqlviz")) + self.assertFalse(_is_module_available("no_such_module_for_test")) + # Missing leaf in an existing (namespace) package. + self.assertFalse(_is_module_available("google.no_such_module_for_test")) + # Missing parent package. + self.assertFalse( + _is_module_available("no_such_module_for_test.submodule") + ) + + def test_module_import_calls_init_extras(self): + import importlib + import google.cloud.managed_spark_connect as msc + + # Configured by class name; not part of the package's API. + self.assertFalse(hasattr(msc, "ManagedSparkConnect")) + self.assertFalse(hasattr(msc, "enable_extras")) + self.assertFalse(hasattr(msc, "disable_extras")) + + # reload() re-executes `from ._ipython import _init_extras`, which + # copies whatever is patched at that moment into the package namespace. + # Undoing the patch does not undo that copy, so reload again outside it + # or every later test in this process sees the mock. + self.addCleanup(importlib.reload, msc) + with mock.patch( + "google.cloud.managed_spark_connect._ipython._init_extras" + ) as mock_init: + importlib.reload(msc) + mock_init.assert_called_once() + + +class TestExtrasInteractiveShell(unittest.TestCase): + + def setUp(self): + import os + from IPython.core.interactiveshell import InteractiveShell + from google.cloud.managed_spark_connect import _ipython + + env_patcher = mock.patch.dict(os.environ) + env_patcher.start() + self.addCleanup(env_patcher.stop) + os.environ.pop(_ipython.ENV_ENABLE_EXTRAS, None) + _ipython._SHELL_STATES.clear() + + self.ip = InteractiveShell.instance() + self.ip.user_ns.pop("explore_dataframe", None) + self.ip.magics_manager.magics.get("line", {}).pop("dpip", None) + self.ip.extension_manager.loaded.discard( + "google.cloud.managed_spark_magics" + ) + if "ManagedSparkConnect" in self.ip.config: + del self.ip.config["ManagedSparkConnect"] + self.ip.configurables[:] = [ + c + for c in self.ip.configurables + if c.__class__.__name__ != "ManagedSparkConnect" + ] + + self._get_ipython_patcher = mock.patch( + "IPython.get_ipython", return_value=self.ip + ) + self._get_ipython_patcher.start() + + def tearDown(self): + from IPython.core.interactiveshell import InteractiveShell + from google.cloud.managed_spark_connect import _ipython + + self._get_ipython_patcher.stop() + self.ip.user_ns.pop("explore_dataframe", None) + self.ip.magics_manager.magics.get("line", {}).pop("dpip", None) + self.ip.extension_manager.loaded.discard( + "google.cloud.managed_spark_magics" + ) + if "ManagedSparkConnect" in self.ip.config: + del self.ip.config["ManagedSparkConnect"] + self.ip.configurables[:] = [ + c + for c in self.ip.configurables + if c.__class__.__name__ != "ManagedSparkConnect" + ] + InteractiveShell.clear_instance() + _ipython._SHELL_STATES.clear() + + def test_real_shell_loads_all_extras_and_toggles_via_config(self): + import io + from google.cloud.managed_spark_connect._ipython import _init_extras + + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + + self.assertIn("explore_dataframe", self.ip.user_ns) + self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = False" + ) + + self.assertNotIn("explore_dataframe", self.ip.user_ns) + self.assertNotIn("dpip", self.ip.magics_manager.magics["line"]) + + with mock.patch("sys.stdout", io.StringIO()): + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = True" + ) + + self.assertIn("explore_dataframe", self.ip.user_ns) + self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + + def test_real_shell_config_magic_pre_and_post_import(self): + import io + from google.cloud.managed_spark_connect._ipython import _init_extras + + # 1. Opt out via %config BEFORE _init_extras() + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = False" + ) + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + + self.assertNotIn("explore_dataframe", self.ip.user_ns) + self.assertNotIn("dpip", self.ip.magics_manager.magics["line"]) + + # 2. Enable via %config AFTER _init_extras() + with mock.patch("sys.stdout", io.StringIO()): + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = True" + ) + + self.assertIn("explore_dataframe", self.ip.user_ns) + self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + + # 3. Disable via %config AFTER _init_extras() + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = False" + ) + + self.assertNotIn("explore_dataframe", self.ip.user_ns) + self.assertNotIn("dpip", self.ip.magics_manager.magics["line"]) + + def test_real_shell_partial_undo_preserves_preloaded_extension(self): + import io + from google.cloud.managed_spark_connect._ipython import _init_extras + + # User manually loads %dpip before managed_spark_connect initializes extras + self.ip.extension_manager.load_extension( + "google.cloud.managed_spark_magics" + ) + self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + self.assertIn("explore_dataframe", self.ip.user_ns) + self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + # %dpip was the user's doing, so the summary must not claim it. + output = stdout_capture.getvalue() + self.assertIn("explore_dataframe()", output) + self.assertNotIn("%dpip", output) + + # Disabling extras via %config must keep %dpip (pre-loaded by user) while removing explore_dataframe + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = False" + ) + + self.assertNotIn("explore_dataframe", self.ip.user_ns) + self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + + def test_real_shell_config_magic_overrides_env_opt_out(self): + import io + import os + from google.cloud.managed_spark_connect._ipython import ( + ENV_ENABLE_EXTRAS, + _init_extras, + ) + + # %config is evaluated before the import, so the merged shell config + # must beat the environment variable rather than being vetoed by it. + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = True" + ) + with mock.patch.dict(os.environ, {ENV_ENABLE_EXTRAS: "false"}): + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + + self.assertIn("explore_dataframe", self.ip.user_ns) + self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + + def test_real_shell_env_opt_out_then_enable_via_config_magic(self): + import io + import os + from google.cloud.managed_spark_connect._ipython import ( + ENV_ENABLE_EXTRAS, + _init_extras, + ) + + with mock.patch.dict(os.environ, {ENV_ENABLE_EXTRAS: "false"}): + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + self.assertNotIn("explore_dataframe", self.ip.user_ns) + + # Turning the trait on must actually load, and the trait must + # report the same thing the shell is really in. + with mock.patch("sys.stdout", io.StringIO()): + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = True" + ) + + self.assertIn("explore_dataframe", self.ip.user_ns) + self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + + +_IMPORT_ORDER_PROBE = """ +import io, sys +from contextlib import redirect_stdout +from IPython.core.interactiveshell import InteractiveShell + +ip = InteractiveShell.instance() +with redirect_stdout(io.StringIO()): + if sys.argv[1] == "magics_first": + import google.cloud.managed_spark_magics + elif sys.argv[1] == "load_ext_first": + ip.run_line_magic("load_ext", "google.cloud.managed_spark_magics") + import google.cloud.managed_spark_connect + +print("explore_dataframe" in ip.user_ns) +print("dpip" in ip.magics_manager.magics["line"]) +""" + + +class TestExtrasImportOrder(unittest.TestCase): + """Each import order must end up with the same set of extras loaded. + + Has to run out of process: once a module is in sys.modules, the import + order that produced it cannot be replayed. + """ + + def _run(self, order): + import os + import subprocess + import sys + + env = dict(os.environ) + env.pop("MANAGED_SPARK_CONNECT_ENABLE_EXTRAS", None) + result = subprocess.run( + [sys.executable, "-c", _IMPORT_ORDER_PROBE, order], + capture_output=True, + text=True, + env=env, + ) + self.assertEqual(result.returncode, 0, result.stderr) + return result.stdout.split() + + def test_all_import_orders_load_all_extras(self): + for order in ("connect_first", "magics_first", "load_ext_first"): + with self.subTest(order=order): + # Importing managed_spark_magics first used to re-enter + # _init_extras() while that module was still initializing, + # which silently dropped %dpip. + self.assertEqual( + self._run(order), + ["True", "True"], + "expected explore_dataframe and %dpip", + ) + + +_NO_INTERACTIVE_DEPS_PROBE = """ +import sys + +# Simulate an install without the [interactive] extra. A None entry makes both +# importlib.util.find_spec() and import treat the module as absent. +for name in ("IPython", "traitlets", "google.colabsqlviz"): + sys.modules[name] = None + +import google.cloud.managed_spark_connect as msc + +print(msc.ManagedSparkSession is not None) +print(hasattr(msc, "ManagedSparkConnect")) +print("google.cloud.managed_spark_connect._ipython" in sys.modules) +""" + + +class TestWithoutInteractiveExtra(unittest.TestCase): + """The base package must not need anything from the [interactive] extra. + + Runs out of process so that nothing already imported by the test runner + (IPython in particular) leaks in. + """ + + def test_import_without_interactive_deps(self): + import os + import subprocess + import sys + + env = dict(os.environ) + env.pop("MANAGED_SPARK_CONNECT_ENABLE_EXTRAS", None) + result = subprocess.run( + [sys.executable, "-c", _NO_INTERACTIVE_DEPS_PROBE], + capture_output=True, + text=True, + env=env, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(result.stdout.split(), ["True", "False", "False"]) + + if __name__ == "__main__": unittest.main() From a223b7df8180376d570f6eb93f7547d9509aad41 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Tue, 6 Oct 2026 11:02:33 -0700 Subject: [PATCH 2/4] feat: add a built-in sparksql magic for streamlined SQL queries Running Spark SQL from a notebook has so far meant installing the third-party sparksql-magic package, which predates Spark Connect and offers no safe way to pass Python values into a query. This adds a clean-room %%sparksql to managed_spark_magics, written only against public pyspark APIs, that works with any SparkSession, classic or Spark Connect. %%sparksql [-s SESSION] [-l LIMIT] [-v VIEW] [-c] [-e] [-q] [--explain [MODE]] [variable] df = %sparksql [-s SESSION] [-v VIEW] [-c] [-e] [--explain [MODE]] SELECT ... The cell form displays the result of the last statement (and optionally stores the DataFrame in variable); the line form is always --quiet and returns the DataFrame directly. Options and the SparkSql config name match the old package's, so existing %config SparkSql.limit lines keep working. The default limit stays 20 rows. - Parameters: :name binds a notebook variable as a typed SQL literal via spark.sql(args=...); IDENTIFIER(:name) covers table and column names. {name} references a Spark DataFrame to query as a table, without allowing arbitrary Python expressions or non-DataFrame interpolation. Both are found with sqlparse, so text inside strings and comments is left alone. - Multiple statements: a cell is split with sqlparse and run in order, only the last result is shown or stored. Every statement's parameters are bound before any of them runs, so a missing variable fails the cell without side effects. A cell starting with BEGIN is passed through whole as a SQL script. - Rendering: values are formatted by Spark with CAST(... AS STRING), so they look the way Spark prints them rather than however they convert to Python (toPandas() turns arrays with nulls into floats and shows binary as b'...'). limit+1 rows are fetched to report truncation exactly. The HTML table shows column types, right-aligns numbers and marks NULLs, with a show()-style text/plain fallback. DDL shows OK. - Errors: a PySparkException is shown as Spark's condition, message and SQLSTATE, which statement failed, and a caret under the offending SQL in the cell, without the Python and JVM stack traces. The caret is left out when Spark's position is not into the cell's own text: after {name} rewriting, or for an error inside a view. %config SparkSql.full_traceback = True re-raises the original. - Session: -s NAME, else the active session, else the only SparkSession in the notebook. The magic is structured as a self-contained IPython extension in managed_spark_magics.sparksql (with its own SparkSqlMagics class and load_ipython_extension / unload_ipython_extension hooks) rather than being bundled onto ManagedSparkMagics. managed_spark_magics delegates to it, so %load_ext google.cloud.managed_spark_magics continues to load both %dpip and %%sparksql, while %load_ext google.cloud.managed_spark_magics.sparksql loads only %%sparksql. This keeps the implementation decoupled from managed_spark_connect so it can eventually be split out into a standalone project, leaving a thin re-export shim behind. sqlparse is a new dependency of the [interactive] extra. It and pyspark are imported lazily, so loading the extension, and %dpip, needs neither. sparksql-magic is dropped from the dev requirements, along with the README and DEVELOPING sections about it. Since managed_spark_magics now provides more than %dpip, the interactive extras loader tracks every magic an extension adds. A same-named magic that was already registered (e.g. %%sparksql from sparksql-magic) is restored after loading and survives unloading, and the summary lists a line+cell magic once, as %%sparksql. tests/local runs the magic end to end against local classic and Spark Connect sessions, and CI now runs it using the runner's preinstalled JDK 17. --- .github/workflows/tests.yaml | 8 +- DEVELOPING.md | 39 +- README.md | 122 +- .../cloud/managed_spark_connect/_ipython.py | 106 +- google/cloud/managed_spark_magics/__init__.py | 19 +- google/cloud/managed_spark_magics/sparksql.py | 1055 +++++++++++++++++ requirements-dev.txt | 2 +- setup.py | 2 + tests/integration/test_session.py | 54 +- tests/local/__init__.py | 1 + tests/local/test_sparksql_local.py | 276 +++++ .../managed_spark_magics/test_sparksql.py | 907 ++++++++++++++ tests/unit/test_init.py | 153 ++- 13 files changed, 2592 insertions(+), 152 deletions(-) create mode 100644 google/cloud/managed_spark_magics/sparksql.py create mode 100644 tests/local/__init__.py create mode 100644 tests/local/test_sparksql_local.py create mode 100644 tests/unit/managed_spark_magics/test_sparksql.py diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index abaa39f5..d298beb5 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -49,4 +49,10 @@ jobs: pip install -r requirements-test.txt - name: Run unit tests - run: python -m pytest tests/unit/ -v --tb=short -n auto \ No newline at end of file + run: python -m pytest tests/unit/ -v --tb=short -n auto + + # Starts Spark in-process, so it runs serially, outside of xdist. + - name: Run local Spark tests + run: python -m pytest tests/local/ -v --tb=short + env: + JAVA_HOME: ${{ env.JAVA_HOME_17_X64 }} \ No newline at end of file diff --git a/DEVELOPING.md b/DEVELOPING.md index b4a6df2b..06a1cae3 100644 --- a/DEVELOPING.md +++ b/DEVELOPING.md @@ -39,38 +39,25 @@ env \ pytest --tb=auto -v ``` -## Testing with Magic Support - -To run tests with magic functionality, install the required dependencies manually: - -```sh -pip install . -pip install IPython sparksql-magic -``` - -Then run tests as normal. Any magic-related tests will automatically detect and use the available dependencies. +## Testing the Interactive Extras -## Testing without Magic Support +The dependencies of the interactive extras (`explore_dataframe`, `%dpip`, +`%%sparksql`) live in the `interactive` extra (`pip install '.[interactive]'`), +and `requirements-dev.txt` already includes them. The tests in +`tests/unit/test_init.py` assume `google-colabsqlviz`, `ipython`, `sqlparse` +and `traitlets` are importable and will fail, not skip, if they are not. +Missing dependencies are simulated within the tests themselves, by setting +`sys.modules[name] = None`. -To run tests without the magic dependencies, simply install the base package: +`tests/local` runs `%%sparksql` end to end against local classic and Spark +Connect sessions. It needs no Google Cloud project, only Java 17 or 21 (set +`JAVA_HOME` if your default `java` is another version, or the tests skip), and +starts a JVM in-process, so run it on its own rather than under `xdist`: ```sh -pip install . -pytest +pytest tests/local ``` -Tests that require magic functionality will be automatically skipped if the dependencies are not available. - -## Testing the Interactive Extras - -The dependencies of the interactive extras (`explore_dataframe`, `%dpip`) live -in the `interactive` extra (`pip install '.[interactive]'`), and -`requirements-dev.txt` already includes them. The tests in -`tests/unit/test_init.py` assume `google-colabsqlviz`, `ipython` and -`traitlets` are importable and will fail, not skip, if they are not. Missing -dependencies are simulated within the tests themselves, by setting -`sys.modules[name] = None`. - The integration tests in particular can take a while to run. To speed up the testing cycle, you can run them in parallel. You can do so using the `xdist` plugin by setting the `-n` flag to the number of parallel runners you want to diff --git a/README.md b/README.md index a9cdf3f3..bb5a03a6 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,7 @@ up a few interactive conveniences: [google-colabsqlviz](https://pypi.org/project/google-colabsqlviz/) is injected into your notebook globals. - The `%dpip` line magic is loaded. +- The [`%%sparksql`](#spark-sql-magic-sparksql) magic is loaded. ```python import google.cloud.managed_spark_connect # extras load here @@ -145,7 +146,7 @@ Each extra is only set up if its dependency is importable, so without the `interactive` extra you get whichever of them your environment already happens to provide. On import, a single line lists the extras that were loaded. -The extras won't override anything you've already set up, for example if `explore_dataframe` is already present, it is left alone. +The extras won't override anything you've already set up, for example if `explore_dataframe` is already present, or `%%sparksql` was loaded from somewhere else, it is left alone. #### Opting out @@ -170,51 +171,100 @@ Or at runtime: An explicit IPython setting takes precedence over the environment variable. -### Using Spark SQL Magic Commands (Jupyter Notebooks) +### Spark SQL Magic (`%%sparksql`) -The package supports the [sparksql-magic](https://github.com/cryeo/sparksql-magic) library for executing Spark SQL queries directly in Jupyter notebooks. +`%%sparksql` runs Spark SQL against the active session, classic or Spark +Connect, and displays the result as a table. It is one of the interactive +extras, so with the `interactive` extra installed it is loaded on import; you +can also load it explicitly with `%load_ext google.cloud.managed_spark_magics`. -**Installation**: To use magic commands, install the required dependencies manually: -```bash -pip install google-cloud-spark-connect -pip install IPython sparksql-magic +```python +%%sparksql +SELECT * FROM your_table ``` -1. Load the magic extension: - ```python - %load_ext sparksql_magic - ``` +The cell displays the first 20 rows by default. Change that limit per cell with +`-l N`, or globally for the notebook with: -2. Configure default settings (optional): - ```python - %config SparkSql.limit=20 - ``` +```python +%config SparkSql.limit = 100 +``` -3. Execute SQL queries: - ```python - %%sparksql - SELECT * FROM your_table - ``` +Options go on the first line, optionally naming a variable to store the +resulting DataFrame in. When a variable name is provided, it receives the full, +un-truncated `DataFrame` (the `-l` display limit only affects what is shown on +screen): -4. Advanced usage with options: - ```python - # Cache results and create a view - %%sparksql --cache --view result_view df - SELECT * FROM your_table WHERE condition = true - ``` +```python +%%sparksql --cache --view result_view df +SELECT * FROM your_table WHERE condition = true +``` -Available options: -- `--cache` / `-c`: Cache the DataFrame -- `--eager` / `-e`: Cache with eager loading -- `--view VIEW` / `-v VIEW`: Create a temporary view -- `--limit N` / `-l N`: Override default row display limit -- `variable_name`: Store result in a variable +- `-l N` / `--limit N`: Maximum rows to display (default: `SparkSql.limit`, 20). +- `-v VIEW` / `--view VIEW`: Register or replace a temporary view with this name. +- `-c` / `--cache`: Cache the resulting DataFrame (lazily). +- `-e` / `--eager`: Compute the result immediately via `localCheckpoint(eager=True)`. +- `-q` / `--quiet`: Suppress table display. +- `--explain [MODE]`: Print the query plan instead of displaying results (`simple`, + `extended`, `codegen`, `cost`, or `formatted`, the default). +- `-s NAME` / `--session NAME`: Run against the SparkSession in variable `NAME`. + Otherwise the active session is used, or else the only SparkSession in the notebook. -See [sparksql-magic](https://github.com/cryeo/sparksql-magic) for more examples. +To return the DataFrame directly for assignment or method chaining instead of +displaying a table, use the line magic form `%sparksql` (which is always quiet): -**Note**: Magic commands are optional. If you only need basic ManagedSparkSession functionality without Jupyter magic support, install only the base package: -```bash -pip install google-cloud-spark-connect +```python +df = %sparksql SELECT * FROM your_table +``` + +#### Passing Python values + +Notebook variables are passed into queries using either `:name` or `{name}`: + +- **`:name` (SQL literals)**: Binds the notebook variable `name` as a typed SQL + parameter marker (via `spark.sql(args=...)`), passed directly to Spark + without string interpolation or SQL injection risk. +- **`IDENTIFIER(:name)` (Catalog & column identifiers)**: In SQL grammar, + parameter markers only substitute for literal values, not table or column + identifiers. Use Spark's built-in `IDENTIFIER()` clause to parameterize table, + view, or column names from a string variable: + + ```python + table, min_id = "my_catalog.db.events", 100 + ``` + + ```python + %%sparksql + SELECT * FROM IDENTIFIER(:table) WHERE id > :min_id + ``` + +- **`{name}` (DataFrames)**: References a Spark `DataFrame` notebook variable + to query as a relation: + + ```python + %%sparksql + SELECT country, count(*) FROM {events_df} GROUP BY country + ``` + +`:name` and `{name}` are only recognized in SQL expression positions; text +inside string literals (e.g. `'{not_a_variable}'` or `':id'`) and SQL comments +is preserved verbatim. + +#### Multiple statements and errors + +A cell can hold multiple statements separated by `;`. They run in sequence, +and only the final statement's result is displayed or stored in the target +variable. If an intermediate statement fails, earlier statements have already +executed. A cell starting with `BEGIN` is passed through whole as a Spark SQL +compound script. + +When a statement fails, `%%sparksql` prints a clean summary with Spark's error +condition, message, SQLSTATE, and a caret pointing at the offending line in the +cell, without JVM or Python wrapper tracebacks. To re-raise the original +exception with the full traceback, set: + +```python +%config SparkSql.full_traceback = True ``` ## Migrating from dataproc-spark-connect diff --git a/google/cloud/managed_spark_connect/_ipython.py b/google/cloud/managed_spark_connect/_ipython.py index e43f7079..da7bee88 100644 --- a/google/cloud/managed_spark_connect/_ipython.py +++ b/google/cloud/managed_spark_connect/_ipython.py @@ -17,7 +17,7 @@ import logging import os import sys -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from traitlets import Bool, default, observe from traitlets.config import Config, Configurable @@ -31,9 +31,14 @@ # Module that must be installed for explore_dataframe() to be injected. _COLABSQLVIZ_MODULE = "google.colabsqlviz" -# (extension module, magic kind, magic name). Each extension is only loaded if -# its module is installed. -_EXTRAS_EXTENSIONS = (("google.cloud.managed_spark_magics", "line", "dpip"),) +# (extension module, ((magic kind, magic name), ...)). Each extension is only +# loaded if its module is installed. +_EXTRAS_EXTENSIONS = ( + ( + "google.cloud.managed_spark_magics", + (("line", "dpip"), ("line", "sparksql"), ("cell", "sparksql")), + ), +) def _is_module_available(name: str) -> bool: @@ -71,7 +76,7 @@ class ManagedSparkConnect(Configurable): enable_extras = Bool( help=( "Whether to automatically load the installed interactive extras " - "(explore_dataframe, %dpip)." + "(explore_dataframe, %dpip, %%sparksql)." ), ).tag(config=True) @@ -106,6 +111,9 @@ class _ShellExtrasState: shell: Any injected_explore_dataframe: Any = None loaded_extensions: Set[str] = field(default_factory=set) + # Per extension we loaded, the (kind, name) magics it added. Magics that + # already existed under the same name were put back and are not listed. + added_magics: Dict[str, List[Tuple[str, str]]] = field(default_factory=dict) config_instance: Optional[ManagedSparkConnect] = None @@ -208,8 +216,35 @@ def _print_extras_failure() -> None: ) +def _shell_magics(ip: Any) -> Dict[str, Dict[str, Any]]: + magics_mgr = getattr(ip, "magics_manager", None) + magics = getattr(magics_mgr, "magics", None) + return magics if isinstance(magics, dict) else {} + + +def _present_magics( + ip: Any, magics: Sequence[Tuple[str, str]] +) -> Dict[Tuple[str, str], Any]: + """The subset of `magics` the shell already has, with their callables.""" + shell_magics = _shell_magics(ip) + present = {} + for kind, name in magics: + kind_dict = shell_magics.get(kind) + if isinstance(kind_dict, dict) and name in kind_dict: + present[(kind, name)] = kind_dict[name] + return present + + +def _is_from_module(fn: Any, module: str) -> bool: + """Whether a registered magic was defined in `module` or a submodule.""" + mod = getattr(fn, "__module__", None) + return isinstance(mod, str) and ( + mod == module or mod.startswith(module + ".") + ) + + def _is_extension_loaded( - ip: Any, ext_name: str, magic_kind: str, magic_name: str + ip: Any, ext_name: str, magics: Sequence[Tuple[str, str]] ) -> bool: mod = sys.modules.get(ext_name) if mod is not None and getattr( @@ -222,14 +257,8 @@ def _is_extension_loaded( if isinstance(loaded, (set, dict, list, tuple)) and ext_name in loaded: return True - magics_mgr = getattr(ip, "magics_manager", None) - magics = getattr(magics_mgr, "magics", None) - if isinstance(magics, dict): - kind_dict = magics.get(magic_kind) - if isinstance(kind_dict, dict) and magic_name in kind_dict: - return True - - return False + # Every magic it would add is already there, so it has nothing to offer. + return len(_present_magics(ip, magics)) == len(magics) def _print_extras_loaded(loaded: List[str]) -> None: @@ -252,6 +281,15 @@ def _magic_display_name(magic_kind: str, magic_name: str) -> str: return f"{prefix}{magic_name}" +def _magics_display_names(magics: Sequence[Tuple[str, str]]) -> List[str]: + """Display names, listing a magic that is both line and cell once.""" + return [ + _magic_display_name(kind, name) + for kind, name in magics + if not (kind == "line" and ("cell", name) in magics) + ] + + def _load_extras_for_shell(ip: Any) -> None: state = _get_shell_state(ip) target_name = "explore_dataframe" @@ -309,10 +347,10 @@ def _load_extras_for_shell(ip: Any) -> None: # 2. Load magic extensions ext_mgr = getattr(ip, "extension_manager", None) if ext_mgr is not None and hasattr(ext_mgr, "load_extension"): - for ext_name, magic_kind, magic_name in _EXTRAS_EXTENSIONS: + for ext_name, magics in _EXTRAS_EXTENSIONS: try: if ext_name in state.loaded_extensions or ( - _is_extension_loaded(ip, ext_name, magic_kind, magic_name) + _is_extension_loaded(ip, ext_name, magics) ): continue if not _is_module_available(ext_name): @@ -321,10 +359,18 @@ def _load_extras_for_shell(ip: Any) -> None: ext_name, ) continue + # A magic of the same name from somewhere else (e.g. the + # sparksql-magic package) must survive our extension loading. + existing = _present_magics(ip, magics) res = ext_mgr.load_extension(ext_name) if res is None: + shell_magics = _shell_magics(ip) + for (kind, name), fn in existing.items(): + shell_magics.setdefault(kind, {})[name] = fn + added = [m for m in magics if m not in existing] state.loaded_extensions.add(ext_name) - loaded.append(_magic_display_name(magic_kind, magic_name)) + state.added_magics[ext_name] = added + loaded.extend(_magics_display_names(added)) except Exception: logger.debug( "Failed to load IPython extension %s", @@ -362,12 +408,16 @@ def _unload_extras_for_shell(ip: Any) -> None: # 2. Unload only the extensions that we loaded ext_mgr = getattr(ip, "extension_manager", None) - magics_mgr = getattr(ip, "magics_manager", None) - magics = getattr(magics_mgr, "magics", None) loaded = getattr(ext_mgr, "loaded", None) - for ext_name, magic_kind, magic_name in _EXTRAS_EXTENSIONS: + for ext_name, magics in _EXTRAS_EXTENSIONS: if ext_name in state.loaded_extensions: + added = state.added_magics.pop(ext_name, ()) + # Magics from this extension that were already present before we + # loaded (e.g. from `%load_ext ...sparksql`) must survive unload. + preserved = _present_magics( + ip, [m for m in magics if m not in added] + ) try: if ext_mgr is not None and hasattr(ext_mgr, "unload_extension"): ext_mgr.unload_extension(ext_name) @@ -378,16 +428,20 @@ def _unload_extras_for_shell(ip: Any) -> None: exc_info=True, ) try: - if isinstance(magics, dict): - kind_dict = magics.get(magic_kind) - if isinstance(kind_dict, dict): - kind_dict.pop(magic_name, None) + shell_magics = _shell_magics(ip) + for (kind, name), fn in preserved.items(): + shell_magics.setdefault(kind, {})[name] = fn + for kind, name in added: + kind_dict = shell_magics.get(kind) + if isinstance(kind_dict, dict) and _is_from_module( + kind_dict.get(name), ext_name + ): + del kind_dict[name] if isinstance(loaded, set): loaded.discard(ext_name) except Exception: logger.debug( - "Failed to clean up magic %s for extension %s", - magic_name, + "Failed to clean up magics for extension %s", ext_name, exc_info=True, ) diff --git a/google/cloud/managed_spark_magics/__init__.py b/google/cloud/managed_spark_magics/__init__.py index eb667258..4a19cbaf 100644 --- a/google/cloud/managed_spark_magics/__init__.py +++ b/google/cloud/managed_spark_magics/__init__.py @@ -12,11 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from . import sparksql from .magics import ManagedSparkMagics def load_ipython_extension(ipython): ipython.register_magics(ManagedSparkMagics) + # %%sparksql is implemented as a standalone extension in `.sparksql` so it + # can be cleanly split out into a separate project later; delegate to it + # here so `%load_ext google.cloud.managed_spark_magics` loads both. + sparksql.load_ipython_extension(ipython) def unload_ipython_extension(ipython): @@ -25,4 +30,16 @@ def unload_ipython_extension(ipython): if isinstance(magics, dict): line_magics = magics.get("line") if isinstance(line_magics, dict): - line_magics.pop("dpip", None) + owner = getattr(line_magics.get("dpip"), "__self__", None) + if isinstance(owner, ManagedSparkMagics): + del line_magics["dpip"] + # If %%sparksql was also loaded as an extension in its own right, it stays + # until that is unloaded. Removing it here would leave IPython believing + # it is still loaded, so a later `%load_ext` of it would do nothing. + ext_mgr = getattr(ipython, "extension_manager", None) + loaded = getattr(ext_mgr, "loaded", None) + if not ( + isinstance(loaded, (set, dict, list, tuple)) + and sparksql.__name__ in loaded + ): + sparksql.unload_ipython_extension(ipython) diff --git a/google/cloud/managed_spark_magics/sparksql.py b/google/cloud/managed_spark_magics/sparksql.py new file mode 100644 index 00000000..080b7a0d --- /dev/null +++ b/google/cloud/managed_spark_magics/sparksql.py @@ -0,0 +1,1055 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Implementation of the %sparksql / %%sparksql magic. + +Everything here works with any pyspark SparkSession, classic or Spark Connect, +and only uses public pyspark APIs. pyspark and sqlparse are imported lazily so +that loading the extension (and with it %dpip) needs neither. + +This module is structured as a self-contained IPython extension (with its own +SparkSqlMagics class and load/unload_ipython_extension entry points) so it can +be cleanly split out into a separate project later, leaving a thin re-export +shim at this module path. +""" + +import dataclasses +import html +import re +import uuid +from typing import ( + Any, + Dict, + List, + Mapping, + NoReturn, + Optional, + Sequence, + Tuple, +) + +from IPython.core.error import UsageError +from IPython.core.magic import ( + Magics, + line_cell_magic, + magics_class, + no_var_expand, +) +from traitlets import Bool, Int, TraitError, validate +from traitlets.config import Configurable + +USAGE = """\ +%%sparksql [-s SESSION] [-l LIMIT] [-v VIEW] [-c] [-e] [-q] [--explain [MODE]] [variable] + + +df = %sparksql [-s SESSION] [-v VIEW] [-c] [-e] [--explain [MODE]] +""" + +EXPLAIN_MODES = ("simple", "extended", "codegen", "cost", "formatted") + +_IDENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") +_PLACEHOLDER_RE = re.compile(r":([A-Za-z_][A-Za-z0-9_]*)\Z") + +# Rough width at which a value wraps in the HTML table, and is cut off in the +# plain-text one. +_MAX_TEXT_WIDTH = 40 + + +class SparkSql(Configurable): + """%config options for %%sparksql. + + Named SparkSql so that `%config SparkSql.limit = N` lines written for the + third-party sparksql-magic package keep working. + """ + + limit = Int( + 20, help="Default number of rows %%sparksql displays (-l overrides)." + ).tag(config=True) + full_traceback = Bool( + False, + help=( + "Show Spark's full exception and traceback for failed queries" + " instead of a short summary." + ), + ).tag(config=True) + + @validate("limit") + def _validate_limit(self, proposal: Mapping[str, Any]) -> Optional[int]: + # Same rule as -l, which parse_options enforces. + val = proposal["value"] + if val is not None and val < 1: + raise TraitError("SparkSql.limit must be at least 1") + return val + + +def _is_sparksql_config(obj: Any) -> bool: + """Compares by module+qualname so that a reloaded SparkSql still matches.""" + cls = type(obj) + return ( + cls.__module__ == SparkSql.__module__ + and cls.__qualname__ == SparkSql.__qualname__ + ) + + +def get_config(shell: Any) -> SparkSql: + """Returns the shell's SparkSql config, registering one if needed. + + Registering it in shell.configurables is what makes `%config SparkSql.x` + take effect immediately rather than only for configurables created later. + """ + for c in getattr(shell, "configurables", ()): + if _is_sparksql_config(c): + return c + # Without a Configurable parent there is no shell config to inherit. An + # invalid value in the shell's config still raises TraitError. + cfg = SparkSql(parent=shell if isinstance(shell, Configurable) else None) + configurables = getattr(shell, "configurables", None) + if isinstance(configurables, list): + configurables.append(cfg) + return cfg + + +# --------------------------------------------------------------------------- +# Option parsing + + +@dataclasses.dataclass +class Options: + session: Optional[str] = None + limit: Optional[int] = None + view: Optional[str] = None + cache: bool = False + eager: bool = False + quiet: bool = False + explain: Optional[str] = None + variable: Optional[str] = None + + +_VALUE_OPTS = { + "-s": "session", + "--session": "session", + "-l": "limit", + "--limit": "limit", + "-v": "view", + "--view": "view", +} +_FLAG_OPTS = { + "-c": "cache", + "--cache": "cache", + "-e": "eager", + "--eager": "eager", + "-q": "quiet", + "--quiet": "quiet", +} +_SHORT_FLAGS = {k[1]: v for k, v in _FLAG_OPTS.items() if len(k) == 2} + + +def _usage_error(msg: str) -> UsageError: + return UsageError(f"{msg}\n\nUsage:\n{USAGE}") + + +def parse_options(line: str, line_form: bool) -> Tuple[Options, str]: + """Parses magic options. + + For the cell form, `line` holds only options and an optional variable name, + and the returned rest is always empty. For the line form, options end at + the first token that is not an option, and everything from there on is + returned verbatim as the SQL. + + This is hand-rolled rather than built on argparse or + IPython.core.magic_arguments for two reasons: + + - The line form's SQL has to come back exactly as written. Using those + means first splitting the line with shlex (as magic_arguments does), + which strips quotes (`'it''s'` comes back as `its`) and loses the + spacing and character offsets needed to slice the SQL out of `line` + verbatim. + - `--explain [MODE]` takes an optional value. With argparse's + `nargs='?'` it would greedily take the next token, i.e. the + `[variable]` name or the first word of the SQL, as the mode. Here the + next token is only taken if it is an explain mode. + """ + opts = Options() + tokens = [(m.group(), m.start()) for m in re.finditer(r"\S+", line)] + i = 0 + rest = "" + + def take_value(flag: str) -> str: + nonlocal i + if i + 1 >= len(tokens): + raise _usage_error(f"{flag} needs a value") + i += 1 + return tokens[i][0] + + while i < len(tokens): + tok, pos = tokens[i] + if not tok.startswith("-") or tok == "-": + if line_form: + rest = line[pos:] + break + if opts.variable is not None: + raise _usage_error(f"Unexpected argument: {tok}") + opts.variable = tok + i += 1 + continue + + flag, eq, inline_value = tok.partition("=") + if flag == "--explain": + if eq: + mode = inline_value + elif ( + i + 1 < len(tokens) + and tokens[i + 1][0].lower() in EXPLAIN_MODES + ): + i += 1 + mode = tokens[i][0] + else: + mode = "formatted" + if mode.lower() not in EXPLAIN_MODES: + raise _usage_error( + f"Unknown explain mode {mode!r}; expected one of" + f" {', '.join(EXPLAIN_MODES)}" + ) + opts.explain = mode.lower() + elif flag in _VALUE_OPTS: + if eq and not inline_value: + raise _usage_error(f"{flag} needs a value") + value = inline_value if eq else take_value(flag) + setattr(opts, _VALUE_OPTS[flag], value) + elif flag in _FLAG_OPTS and not eq: + setattr(opts, _FLAG_OPTS[flag], True) + elif ( + not flag.startswith("--") + and len(flag) > 2 + and flag[1] in "slv" + and not eq + ): + # Attached short value, e.g. -l50. + setattr(opts, _VALUE_OPTS[flag[:2]], flag[2:]) + elif ( + not flag.startswith("--") + and not eq + and all(ch in _SHORT_FLAGS for ch in flag[1:]) + ): + # Combined short flags, e.g. -ce. + for ch in flag[1:]: + setattr(opts, _SHORT_FLAGS[ch], True) + else: + raise _usage_error(f"Unknown option: {tok}") + i += 1 + + if opts.limit is not None: + try: + limit = int(opts.limit) + except ValueError: + raise _usage_error( + f"--limit must be an integer, got {opts.limit!r}" + ) + if limit < 1: + raise _usage_error("--limit must be at least 1") + opts.limit = limit + for attr in ("session", "variable"): + value = getattr(opts, attr) + if value is not None and not _IDENT_RE.match(value): + raise _usage_error(f"Not a valid Python identifier: {value!r}") + if opts.view is not None and not _IDENT_RE.match(opts.view): + # Temporary views are unqualified; catching a bad name here beats a + # Spark error after the query has already run. + raise _usage_error(f"Not a valid temporary view name: {opts.view!r}") + if line_form: + if opts.limit is not None: + # The line form never displays anything, so this would be + # silently ignored. + raise _usage_error("-l only applies to the cell form (%%sparksql)") + if not rest.strip(): + raise _usage_error("No SQL given") + opts.quiet = True + return opts, rest + + +# --------------------------------------------------------------------------- +# Statement splitting and parameter binding + + +def _import_sqlparse() -> Any: + try: + import sqlparse + except ImportError: + raise UsageError( + "%%sparksql needs the sqlparse package. Install it with:\n" + " pip install 'google-cloud-spark-connect[interactive]'" + ) from None + return sqlparse + + +@dataclasses.dataclass +class Statement: + """One SQL statement from a cell, plus where it came from.""" + + text: str + # Offset of `text` within the cell, for mapping error positions back. + offset: int + # Flattened sqlparse tokens making up exactly `text`, as (ttype, value). + tokens: List[Tuple[Any, str]] + + +def _is_code(token: Tuple[Any, str]) -> bool: + """Whether a token is significant: not whitespace, a comment or `;`.""" + from sqlparse import tokens as T + + ttype, value = token + return ( + value.strip() not in ("", ";") + and ttype not in T.Comment + and ttype not in T.Text.Whitespace + ) + + +def _make_statement( + flat: List[Tuple[Any, str]], start: int +) -> Optional[Statement]: + """Trims leading whitespace and everything after the last code token.""" + code = [i for i, tok in enumerate(flat) if _is_code(tok)] + if not code: + return None + first = next(i for i, (_, v) in enumerate(flat) if v.strip()) + kept = flat[first : code[-1] + 1] + offset = start + sum(len(v) for _, v in flat[:first]) + return Statement("".join(v for _, v in kept), offset, kept) + + +def split_statements(sql: str) -> List[Statement]: + """Splits a cell into statements, keeping each one's offset in the cell. + + A cell whose first keyword is BEGIN is a SQL script (a single compound + statement), and is never split. + """ + sqlparse = _import_sqlparse() + + parsed = [] + offset = 0 + for stmt in sqlparse.parse(sql): + raw = str(stmt) + parsed.append((offset, [(t.ttype, t.value) for t in stmt.flatten()])) + offset += len(raw) + + all_tokens = [t for _, flat in parsed for t in flat] + first_code = next((v for tt, v in all_tokens if _is_code((tt, v))), None) + if first_code is None: + return [] + if first_code.upper() == "BEGIN": + return [_make_statement(all_tokens, 0)] + statements = (_make_statement(flat, start) for start, flat in parsed) + return [s for s in statements if s is not None] + + +def _brace_ident( + toks: Sequence[Tuple[Any, str]], i: int +) -> Optional[Tuple[int, str]]: + """If `toks[i]` opens a `{name}` reference, returns (closing_index, name). + + Matches on token values alone, not types: sqlparse has no rule for braces + and emits each as a one-character Error token, but that is incidental, + and its lexer is process-wide and extensible. A brace inside a string + literal, quoted name or comment is part of that longer token, so it never + makes up a token's whole value. + """ + if toks[i][1] != "{": + return None + j = i + 1 + while j < len(toks) and toks[j][1] not in ("{", "}"): + j += 1 + if j < len(toks) and toks[j][1] == "}": + inner = "".join(v for _, v in toks[i + 1 : j]) + if _IDENT_RE.match(inner): + return j, inner + return None + + +def find_parameters(statement: Statement) -> Tuple[List[str], List[str]]: + """Returns the (`:name` marker names, `{name}` DataFrame names) in a statement. + + Both are only recognized outside string literals and comments. + """ + from sqlparse import tokens as T + + markers: List[str] = [] + fields: List[str] = [] + toks = statement.tokens + i = 0 + while i < len(toks): + ttype, value = toks[i] + if ttype in T.Name.Placeholder: + m = _PLACEHOLDER_RE.match(value) + if m and m.group(1) not in markers: + markers.append(m.group(1)) + else: + brace = _brace_ident(toks, i) + if brace is not None: + j, name = brace + if name not in fields: + fields.append(name) + i = j + i += 1 + return markers, fields + + +def escape_braces(statement: Statement) -> str: + """Returns the statement text with braces escaped for str.format(). + + Spark formats the whole query with str.format() when it is given keyword + arguments, so every brace that is not part of a `{name}` reference, such + as one inside a string literal or a non-identifier `{...}`, has to be + doubled first. + """ + out = [] + toks = statement.tokens + i = 0 + while i < len(toks): + brace = _brace_ident(toks, i) + if brace is not None: + j, name = brace + out.append(f"{{{name}}}") + i = j + 1 + else: + _, value = toks[i] + out.append(value.replace("{", "{{").replace("}", "}}")) + i += 1 + return "".join(out) + + +def _dataframe_classes() -> Tuple[type, ...]: + classes = [] + try: + from pyspark.sql import DataFrame + + classes.append(DataFrame) + except ImportError: + pass + try: + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + + if ConnectDataFrame not in classes: + classes.append(ConnectDataFrame) + except ImportError: + pass + return tuple(classes) + + +def _type_name(val: Any) -> str: + cls = type(val) + if cls.__name__ == "DataFrame": + return f"{cls.__module__.split('.')[0]}.DataFrame" + return cls.__name__ + + +def bind_parameters( + statement: Statement, user_ns: Mapping[str, Any] +) -> Tuple[str, Dict[str, Any], Dict[str, Any]]: + """Returns (sql, args, kwargs) to pass to spark.sql() for a statement.""" + markers, fields = find_parameters(statement) + missing = [n for n in markers + fields if n not in user_ns] + if missing: + names = ", ".join(dict.fromkeys(missing)) + raise UsageError( + f"Not defined in the notebook: {names}. SQL parameters (:name) and" + " DataFrame references ({name}) are bound from notebook variables." + ) + df_classes = _dataframe_classes() if (markers or fields) else () + for n in markers: + if isinstance(user_ns[n], df_classes): + raise UsageError( + f":{n}: expected a SQL value, got a Spark DataFrame." + f" Use {{{n}}} to query a DataFrame." + ) + for n in fields: + val = user_ns[n] + if not isinstance(val, df_classes): + raise UsageError( + f"{{{n}}}: expected a Spark DataFrame, got {_type_name(val)}." + f" Use :{n} for SQL values, or IDENTIFIER(:{n}) for table and" + " column names." + ) + args = {n: user_ns[n] for n in markers} + kwargs = {n: user_ns[n] for n in fields} + sql = escape_braces(statement) if kwargs else statement.text + return sql, args, kwargs + + +# --------------------------------------------------------------------------- +# Session resolution + + +def _session_classes() -> Tuple[type, ...]: + classes = [] + try: + from pyspark.sql import SparkSession + + classes.append(SparkSession) + except ImportError: + pass + try: + from pyspark.sql.connect.session import SparkSession as ConnectSession + + classes.append(ConnectSession) + except ImportError: + pass + return tuple(classes) + + +def _active_session() -> Any: + try: + from pyspark.sql.connect.session import SparkSession as ConnectSession + + return ConnectSession.active() + except Exception: + pass + try: + from pyspark.sql import SparkSession + + return SparkSession.active() + except Exception: + return None + + +def resolve_session(user_ns: Mapping[str, Any], name: Optional[str]) -> Any: + """Finds the SparkSession to run against. + + In order: the -s variable, the active session, and finally the only + SparkSession in the notebook namespace. + """ + classes = _session_classes() + if not classes: + raise UsageError("%%sparksql needs pyspark, which is not installed.") + if name is not None: + if name not in user_ns: + raise UsageError(f"-s {name}: no such variable") + session = user_ns[name] + if not isinstance(session, classes): + raise UsageError( + f"-s {name}: expected a SparkSession, got" + f" {type(session).__name__}" + ) + return session + + session = _active_session() + if session is not None: + return session + + found = [ + k + for k, v in user_ns.items() + if not k.startswith("_") and isinstance(v, classes) + ] + if len(found) == 1: + return user_ns[found[0]] + if found: + raise UsageError( + f"Found several SparkSessions ({', '.join(sorted(found))}) and" + " none is active; pick one with -s NAME." + ) + raise UsageError( + "No SparkSession found. Create one first (for example with" + " ManagedSparkSession.builder.getOrCreate()), or pass -s NAME." + ) + + +# --------------------------------------------------------------------------- +# Errors + + +class SparkSqlError(Exception): + """A failed SQL statement, rendered without the Python traceback. + + The traceback would only show the magic's and pyspark's internals; the + error message and the position in the SQL are what matter. The original + exception is kept as __cause__. + + Not a subclass of managed_spark_connect's exception type: importing that + package pulls in pyspark, which loading this extension must not require. + """ + + def __init__(self, message: str): + super().__init__(message) + self.message = message + + def _render_traceback_(self) -> List[str]: + return self.message.splitlines() + + +def _condition(exc: Any) -> Optional[str]: + # getErrorClass() is the pre-4.0 name, deprecated since. + fn = getattr(exc, "getCondition", None) or getattr( + exc, "getErrorClass", None + ) + if fn is None: + return None + try: + return fn() or None + except Exception: + return None + + +def _sql_context(exc: Any, sql: str) -> Optional[Tuple[int, int]]: + """Returns the (start, stop) offsets in `sql` Spark blamed, if it said. + + A context only counts if its fragment is the text at those offsets in + `sql`: an error raised inside a view, for example, carries offsets into + the view's own SQL, which would put the caret somewhere unrelated. + """ + try: + contexts = exc.getQueryContext() + except Exception: + return None + from pyspark.errors import QueryContextType + + for ctx in contexts or (): + try: + if ctx.contextType() != QueryContextType.SQL: + continue + start, stop = ctx.startIndex(), ctx.stopIndex() + get_fragment = getattr(ctx, "fragment", None) + fragment = get_fragment() if get_fragment is not None else None + except Exception: + continue + if start is None or stop is None or not 0 <= start <= stop: + continue + if fragment and sql[start : stop + 1] != fragment: + continue + return start, stop + return None + + +# The position Spark writes into parse errors, which have no query context: +# "Syntax error at or near 'SELEC'. SQLSTATE: 42601 (line 1, pos 0)". +_PARSE_POSITION_RE = re.compile(r"\(line (\d+), pos (\d+)\)") + + +def _parse_position(exc: Any, sql: str) -> Optional[Tuple[int, int]]: + """Returns the (start, stop) of the token a parse error points at.""" + m = _PARSE_POSITION_RE.search(str(exc)) + if not m: + return None + line, pos = int(m.group(1)), int(m.group(2)) + start = 0 + for _ in range(line - 1): + start = sql.find("\n", start) + 1 + if start == 0: + return None + start += pos + if start >= len(sql): + return None + token = re.match(r"\w+|\S", sql[start:]) + return start, start + (len(token.group()) if token else 1) - 1 + + +# Where Spark's own message ends and the extras it appends begin: the JVM +# stack trace (classic, and Connect's str()), an "== SQL ==" excerpt, and on +# Connect the analyzed plan after "; line 2 pos 2;" or a parse error's +# " (line 1, pos 0)". +_MESSAGE_END_RE = re.compile( + r"\n\s*\n|\n== SQL|\nJVM stacktrace:|;\s*line \d+ pos \d+;?" + r"| \(line \d+, pos \d+\)" +) + + +def _first_paragraph(exc: Any, condition: Optional[str]) -> str: + """Spark's message, without the context and JVM stack it appends.""" + text = None + get_message = getattr(exc, "getMessage", None) + if get_message is not None: + try: + text = get_message() + except Exception: + text = None + if not text: + text = str(exc) + text = _MESSAGE_END_RE.split(text.strip(), maxsplit=1)[0].rstrip() + if condition and not text.startswith("["): + text = f"[{condition}] {text}" + return text + + +def format_error( + exc: Any, + cell: str, + statement: Statement, + index: int, + count: int, + offsets_match: bool, +) -> str: + """Formats a Spark SQL error from `statement` for display.""" + condition = _condition(exc) + lines = [_first_paragraph(exc, condition)] + sqlstate = None + get_state = getattr(exc, "getSqlState", None) + if get_state is not None: + try: + sqlstate = get_state() + except Exception: + sqlstate = None + if sqlstate and f"SQLSTATE: {sqlstate}" not in lines[0]: + lines[0] += f" SQLSTATE: {sqlstate}" + + where = [] + if count > 1: + where.append(f"statement {index + 1} of {count}") + if index > 0: + where.append("earlier statements already ran") + span = None + if offsets_match: + span = _sql_context(exc, statement.text) or _parse_position( + exc, statement.text + ) + if span is not None and span[0] < len(statement.text): + start = statement.offset + span[0] + stop = statement.offset + min(span[1], len(statement.text) - 1) + line_start = cell.rfind("\n", 0, start) + 1 + line_end = cell.find("\n", start) + if line_end < 0: + line_end = len(cell) + line_no = cell.count("\n", 0, start) + 1 + where.insert(0, f"line {line_no}") + src = cell[line_start:line_end] + width = max(1, min(stop, line_end - 1) - start + 1) + caret = " " * (start - line_start) + "^" * width + excerpt = [f" {src.rstrip()}", f" {caret}"] + else: + excerpt = [] + if where: + lines.append(f"({'; '.join(where)})") + lines.extend(excerpt) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Rendering + +_NUMERIC_TYPES = ( + "ByteType", + "ShortType", + "IntegerType", + "LongType", + "FloatType", + "DoubleType", + "DecimalType", +) + + +@dataclasses.dataclass +class Table: + names: List[str] + types: List[str] + numeric: List[bool] + rows: List[List[Optional[str]]] + truncated: bool + + +def fetch_table(df: Any, limit: int) -> Table: + """Fetches up to `limit` rows, formatted to strings by Spark. + + Formatting server-side with CAST(... AS STRING) means values look the way + Spark itself prints them (timestamps in the session time zone, nested + types, decimals), rather than however they convert to Python. Fetching one + extra row tells us whether there are more without a separate count. + """ + # SQL expressions rather than pyspark.sql.functions: which implementation + # those dispatch to, classic or Connect, depends on process-wide state + # rather than on `df`. + fields = df.schema.fields + # Positional names sidestep duplicate or dotted column names. + tmp = [f"_c{i}" for i in range(len(fields))] + exprs = [] + for name, field in zip(tmp, fields): + expr = name + if field.dataType.typeName() == "binary": + # Same "[62 69 6E]" form df.show() uses. + expr = ( + f"concat('[', regexp_replace(hex({name}), '(..)(?=.)', '$1 ')," + " ']')" + ) + exprs.append(f"CAST({expr} AS STRING) AS {name}") + rows = df.toDF(*tmp).selectExpr(*exprs).limit(limit + 1).collect() + return Table( + names=[f.name for f in fields], + types=[f.dataType.simpleString() for f in fields], + numeric=[type(f.dataType).__name__ in _NUMERIC_TYPES for f in fields], + rows=[list(r) for r in rows[:limit]], + truncated=len(rows) > limit, + ) + + +def _footer(table: Table) -> str: + n = len(table.rows) + if table.truncated: + return f"Showing the first {n} rows; more are available (-l to change)." + return f"{n} row{'' if n == 1 else 's'}" + + +def render_html(table: Table) -> str: + cls = f"sparksql-{uuid.uuid4().hex[:8]}" + out = [ + f'
", + ] + for name, typ in zip(table.names, table.types): + out.append( + f"' + ) + out.append("") + for row in table.rows: + out.append("") + for value, numeric in zip(row, table.numeric): + classes = [] + if numeric: + classes.append("n") + if value is None: + classes.append("null") + value = "NULL" + attr = f' class="{" ".join(classes)}"' if classes else "" + out.append(f"{html.escape(value)}") + out.append("") + out.append("
{html.escape(name)}" + f'{html.escape(typ)}
") + out.append(f'
{html.escape(_footer(table))}
') + return "".join(out) + + +def _cut(value: str) -> str: + value = value.replace("\n", "\\n") + if len(value) > _MAX_TEXT_WIDTH: + return value[: _MAX_TEXT_WIDTH - 3] + "..." + return value + + +def render_text(table: Table) -> str: + """Renders the table the way df.show() would, for terminals.""" + cells = [ + [_cut("NULL" if v is None else v) for v in row] for row in table.rows + ] + widths = [ + max([len(name)] + [len(row[i]) for row in cells]) + for i, name in enumerate(table.names) + ] + sep = "+" + "+".join("-" * w for w in widths) + "+" + + def fmt(values: Sequence[str]) -> str: + return "|" + "|".join(v.rjust(w) for v, w in zip(values, widths)) + "|" + + lines = [sep, fmt(table.names), sep] + lines.extend(fmt(row) for row in cells) + lines.append(sep) + lines.append(_footer(table)) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Running + + +def run(shell: Any, line: str, cell: Optional[str]) -> Any: + """Entry point for both forms of the magic.""" + line_form = cell is None + opts, rest = parse_options(line, line_form) + sql = rest if line_form else cell + + statements = split_statements(sql) + if not statements: + raise _usage_error("No SQL given") + + config = get_config(shell) + user_ns = shell.user_ns + session = resolve_session(user_ns, opts.session) + + # Bind every statement before running any, so that a missing variable in a + # later statement fails the cell before earlier ones have side effects. + bound = [bind_parameters(s, user_ns) for s in statements] + df = None + for index, (statement, (bound_sql, args, kwargs)) in enumerate( + zip(statements, bound) + ): + try: + df = session.sql(bound_sql, args=args or None, **kwargs) + except Exception as e: + _reraise(e, config, sql, statement, index, len(statements), kwargs) + # Errors from here on, e.g. at collect time, still carry positions in the + # last statement's SQL as Spark saw it. + last_kwargs = bound[-1][2] + + try: + if opts.eager: + # The checkpoint is already materialized and stored, so caching + # it as well (-c -e) would only keep a second copy. + df = df.localCheckpoint(eager=True) + elif opts.cache: + df = df.cache() + if opts.view is not None: + df.createOrReplaceTempView(opts.view) + if opts.explain is not None: + df.explain(mode=opts.explain) + elif not opts.quiet: + _display(df, opts.limit or config.limit) + # Only once everything has succeeded, so that a failure does not + # leave the variable bound to a DataFrame that does not work. + if opts.variable is not None: + user_ns[opts.variable] = df + except Exception as e: + _reraise( + e, + config, + sql, + statements[-1], + len(statements) - 1, + len(statements), + last_kwargs, + ) + + return df if line_form else None + + +def _reraise( + exc: Exception, + config: SparkSql, + cell: str, + statement: Statement, + index: int, + count: int, + kwargs: Mapping[str, Any], +) -> NoReturn: + try: + from pyspark.errors import PySparkException + except ImportError: + raise exc + if not isinstance(exc, PySparkException): + raise exc + if config.full_traceback: + if count > 1: + add_note = getattr(exc, "add_note", None) + if add_note is not None: + add_note(f"In %%sparksql statement {index + 1} of {count}") + raise exc + # With {name} references Spark rewrites the SQL before parsing it, so its + # error positions no longer line up with the cell. + message = format_error( + exc, cell, statement, index, count, offsets_match=not kwargs + ) + raise SparkSqlError(message) from exc + + +def _display(df: Any, limit: int) -> None: + from IPython.display import display + + if not df.schema.fields: + # DDL and most DML: there is nothing to show but that it worked. + display({"text/plain": "OK"}, raw=True) + return + # TODO: Route through colabsqlviz's explore_dataframe() once it can be + # asked to open as a table rather than a chart. + table = fetch_table(df, limit) + display( + {"text/html": render_html(table), "text/plain": render_text(table)}, + raw=True, + ) + + +# --------------------------------------------------------------------------- +# IPython extension +# +# Kept in this module (rather than bundled into ManagedSparkMagics) so that +# %%sparksql can be cleanly split out into a standalone project later. + + +@magics_class +class SparkSqlMagics(Magics): + + def __init__(self, shell: Any, **kwargs: Any) -> None: + super().__init__(shell, **kwargs) + # Registered up front so `%config SparkSql.limit = N` works before the + # first %%sparksql cell runs. + get_config(shell) + + @no_var_expand + @line_cell_magic + def sparksql(self, line: str, cell: Optional[str] = None) -> Any: + """Runs Spark SQL against the active SparkSession. + + Cell form: displays the result of the last statement. + + %%sparksql [-s SESSION] [-l LIMIT] [-v VIEW] [-c] [-e] [-q] + [--explain [MODE]] [variable] + SELECT * FROM t WHERE x = :x + + Line form: always --quiet; returns the DataFrame directly. + + df = %sparksql [-s SESSION] [-v VIEW] [-c] [-e] + [--explain [MODE]] SELECT * FROM {other_df} + + Options: + variable Store the un-truncated DataFrame in this variable + (cell form only; -l only limits displayed rows). + -s, --session N Use the SparkSession in variable N. Default: the + active session, else the only one in the notebook. + -l, --limit N Maximum rows to display (default: %config + SparkSql.limit, initially 20). + -v, --view V Also register the result as temp view V. + -c, --cache Cache the result (lazily). Redundant with -e. + -e, --eager Materialize the result now, with + localCheckpoint(eager=True). + -q, --quiet Don't display anything. + --explain [MODE] Print the query plan instead of results. MODE is + simple, extended, codegen, cost or formatted + (default). + + Python values are bound from notebook variables: + :name A typed SQL literal (e.g. WHERE id = :id). For a + table or column name, use IDENTIFIER(:name). + {name} A Spark DataFrame to query as a table (e.g. + SELECT * FROM {df}). + Literals and comments: :name and {name} inside string literals and + comments are preserved verbatim. + + A cell can hold several statements separated by ';'. They run in + order, and only the last one's result is displayed or stored. A cell + starting with BEGIN is run whole, as a SQL script. + """ + return run(self.shell, line, cell) + + +def load_ipython_extension(ipython: Any) -> None: + ipython.register_magics(SparkSqlMagics) + + +def unload_ipython_extension(ipython: Any) -> None: + magics_mgr = getattr(ipython, "magics_manager", None) + magics = getattr(magics_mgr, "magics", None) + if not isinstance(magics, dict): + return + for kind in ("line", "cell"): + kind_magics = magics.get(kind) + if not isinstance(kind_magics, dict): + continue + # Leave alone a magic of the same name that someone else + # registered, e.g. %%sparksql from the sparksql-magic package. + owner = getattr(kind_magics.get("sparksql"), "__self__", None) + if isinstance(owner, SparkSqlMagics): + del kind_magics["sparksql"] diff --git a/requirements-dev.txt b/requirements-dev.txt index 09806396..c4ad48ba 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -7,7 +7,7 @@ packaging>=20.0 pyink~=24.0 pyspark[connect]~=4.0.0 setuptools>=72.0 -sparksql-magic>=0.0.3 +sqlparse>=0.5.0 tqdm>=4.67 traitlets>=5.1 websockets>=14.0 diff --git a/setup.py b/setup.py index b1850daa..0c235823 100644 --- a/setup.py +++ b/setup.py @@ -45,6 +45,8 @@ # managed_spark_magics; previously these only arrived transitively # via google-colabsqlviz. "ipython>=8.0", + # Statement splitting and parameter detection in %%sparksql. + "sqlparse>=0.5.0", "traitlets>=5.1", ], }, diff --git a/tests/integration/test_session.py b/tests/integration/test_session.py index c9217137..b2cf110e 100644 --- a/tests/integration/test_session.py +++ b/tests/integration/test_session.py @@ -587,64 +587,28 @@ def test_session_id_validation_in_integration( assert builder._custom_session_id == valid_id -def test_sparksql_magic_library_available(connect_session): - """Test that sparksql-magic library can be imported and loaded.""" - pytest.importorskip( - "IPython", reason="IPython not available (install with magic extra)" - ) - pytest.importorskip( - "sparksql_magic", - reason="sparksql-magic not available (install with magic extra)", - ) - - from IPython.terminal.interactiveshell import TerminalInteractiveShell - - # Create real IPython shell - shell = TerminalInteractiveShell.instance() - shell.user_ns = {"spark": connect_session} - - # Test that sparksql_magic can be loaded (this verifies the dependency works) - try: - shell.run_line_magic("load_ext", "sparksql_magic") - magic_loaded = True - except Exception as e: - magic_loaded = False - print(f"Failed to load sparksql_magic: {e}") - - assert magic_loaded, "sparksql_magic should be available as a dependency" - - # Test that ManagedSparkSession can execute SQL (ensuring basic compatibility) - result = connect_session.sql("SELECT 'integration_test' as test_column") - data = result.collect() - assert len(data) == 1 - assert data[0]["test_column"] == "integration_test" - - def test_sparksql_magic_with_managed_spark_session(connect_session): - """Test that sparksql-magic works with registered ManagedSparkSession.""" + """Test that %%sparksql works with the active ManagedSparkSession.""" pytest.importorskip( - "IPython", reason="IPython not available (install with magic extra)" + "IPython", reason="IPython not available (install [interactive])" ) pytest.importorskip( - "sparksql_magic", - reason="sparksql-magic not available (install with magic extra)", + "sqlparse", reason="sqlparse not available (install [interactive])" ) from IPython.terminal.interactiveshell import TerminalInteractiveShell - # Create real IPython shell (ManagedSparkSession is already registered globally) + # ManagedSparkSession is already the active session. shell = TerminalInteractiveShell.instance() + shell.run_line_magic("load_ext", "google.cloud.managed_spark_magics") + shell.user_ns["factor"] = 5 - # Load the sparksql_magic extension - shell.run_line_magic("load_ext", "sparksql_magic") - - # Test sparksql magic with SQL expressions (no variable capture to avoid namespace issues) shell.run_cell_magic( "sparksql", - "result_df", + "-q result_df", """ - SELECT - 10 * 5 as multiplication, + SELECT + 10 * :factor as multiplication, SQRT(16) as square_root, CONCAT('Dataproc', '-', 'Spark') as joined_string """, diff --git a/tests/local/__init__.py b/tests/local/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/tests/local/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/local/test_sparksql_local.py b/tests/local/test_sparksql_local.py new file mode 100644 index 00000000..12e79f89 --- /dev/null +++ b/tests/local/test_sparksql_local.py @@ -0,0 +1,276 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""%%sparksql against real local Spark, both classic and Spark Connect. + +These start an in-process JVM, so they need Java 17 or 21 and are kept out of +tests/unit. Run them on their own, without xdist: + + python -m pytest tests/local +""" + +import io +import os +import re +import shutil +import subprocess +from contextlib import redirect_stdout +from unittest import mock + +import pytest +from IPython.core.interactiveshell import InteractiveShell +from pyspark.errors import PySparkException +from pyspark.sql import SparkSession +from pyspark.util import is_remote_only + +from google.cloud.managed_spark_magics import sparksql + + +def _java_version(): + """The major version of the Java pyspark would launch, or None.""" + java_home = os.environ.get("JAVA_HOME") + java = ( + os.path.join(java_home, "bin", "java") + if java_home + else shutil.which("java") + ) + if not java: + return None + try: + out = subprocess.run( + [java, "-version"], capture_output=True, text=True, check=True + ).stderr + except (OSError, subprocess.SubprocessError): + return None + m = re.search(r'version "(\d+)', out) + return int(m.group(1)) if m else None + + +_JAVA_VERSION = _java_version() + +pytestmark = [ + pytest.mark.skipif( + # Spark 4.0 supports exactly these; newer JDKs fail at JVM startup. + _JAVA_VERSION not in (17, 21), + reason=( + f"needs Java 17 or 21 (found {_JAVA_VERSION}); point JAVA_HOME at" + " one" + ), + ), + pytest.mark.skipif( + is_remote_only(), + reason="needs full pyspark, with its jars, not pyspark-client", + ), +] + + +@pytest.fixture(scope="module", params=["classic", "connect"]) +def spark(request): + with pytest.MonkeyPatch.context() as mp: + builder = SparkSession.builder + if request.param == "classic": + builder = builder.master("local[2]").config( + "spark.ui.enabled", "false" + ) + else: + # Makes the local Connect server bind an ephemeral port, and the + # client find it, rather than both assuming 15002, which may be + # taken. + mp.setenv("SPARK_TESTING", "1") + builder = builder.remote("local[2]") + session = ( + builder.appName(f"sparksql-{request.param}") + .config("spark.sql.shuffle.partitions", "2") + .config("spark.sql.scripting.enabled", "true") + .getOrCreate() + ) + yield session + session.stop() + + +@pytest.fixture +def ip(spark): + shell = InteractiveShell.instance() + shell.run_line_magic( + "load_ext", "google.cloud.managed_spark_magics.sparksql" + ) + sparksql.get_config(shell).limit = 20 + sparksql.get_config(shell).full_traceback = False + shell.user_ns["spark"] = spark + yield shell + shell.run_line_magic( + "unload_ext", "google.cloud.managed_spark_magics.sparksql" + ) + InteractiveShell.clear_instance() + + +@pytest.fixture +def displayed(): + """The bundles %%sparksql displayed, in order.""" + calls = [] + with mock.patch( + "IPython.display.display", + side_effect=lambda obj, raw=False: calls.append(obj), + ): + yield calls + + +def test_select_renders_table(ip, displayed): + ip.run_cell_magic( + "sparksql", + "", + "SELECT id, CAST(id AS STRING) AS s FROM range(3) ORDER BY id", + ) + (bundle,) = displayed + html = bundle["text/html"] + assert 'bigint' in html + assert 'string' in html + assert '2' in html + assert "3 rows" in html + assert bundle["text/plain"].splitlines()[1] == "|id|s|" + + +def test_limit_and_truncation(ip, displayed): + ip.run_cell_magic("sparksql", "-l 2", "SELECT * FROM range(5)") + assert "Showing the first 2 rows" in displayed[0]["text/html"] + sparksql.get_config(ip).limit = 5 + ip.run_cell_magic("sparksql", "", "SELECT * FROM range(5)") + assert "5 rows" in displayed[1]["text/html"] + + +def test_values_are_formatted_by_spark(ip, displayed): + ip.run_cell_magic( + "sparksql", + "", + "SELECT CAST('bin' AS BINARY) AS b, CAST(NULL AS INT) AS n," + " array(1, NULL) AS a, map('k', 1.5D) AS m, DATE'2026-01-02' AS d", + ) + html = displayed[0]["text/html"] + assert "[62 69 6E]" in html + assert 'NULL' in html + assert "[1, null]" in html + assert "{k -> 1.5}" in html + assert "2026-01-02" in html + + +def test_ddl_displays_ok(ip, displayed): + ip.run_cell_magic( + "sparksql", "", "CREATE OR REPLACE TEMP VIEW ddl_v AS SELECT 1 AS a" + ) + assert displayed == [{"text/plain": "OK"}] + + +def test_parameter_markers(ip, spark, displayed): + spark.range(10).createOrReplaceTempView("param_t") + ip.user_ns.update(t="param_t", lo=3, hi=5) + ip.run_cell_magic( + "sparksql", + "out", + "SELECT id FROM IDENTIFIER(:t) WHERE id BETWEEN :lo AND :hi" + " AND ':nope' <> '' ORDER BY id", + ) + assert [r.id for r in ip.user_ns["out"].collect()] == [3, 4, 5] + + +def test_dataframe_references(ip, spark): + from IPython.core.error import UsageError + + ip.user_ns["df"] = spark.range(4) + ip.user_ns["label"] = "it's" + df = ip.run_line_magic( + "sparksql", + "SELECT :label AS l, '{not a field}' AS s, count(*) AS c FROM {df}", + ) + assert df.collect()[0].asDict() == { + "l": "it's", + "s": "{not a field}", + "c": 4, + } + with pytest.raises( + UsageError, match=r"\{label\}: expected a Spark DataFrame" + ): + ip.run_line_magic("sparksql", "SELECT {label}") + + +def test_multiple_statements(ip, displayed): + ip.run_cell_magic( + "sparksql", + "out", + "CREATE OR REPLACE TEMP VIEW multi_v AS SELECT 1 AS a;\n" + "-- a comment\n" + "SELECT a + 1 AS b FROM multi_v;\n", + ) + assert ip.user_ns["out"].collect()[0].b == 2 + assert len(displayed) == 1 + + +def test_script(ip, displayed): + ip.run_cell_magic( + "sparksql", + "out", + "BEGIN\n DECLARE x INT = 41;\n SELECT x + 1 AS y;\nEND", + ) + assert ip.user_ns["out"].collect()[0].y == 42 + + +def test_view_cache_eager(ip, spark, displayed): + ip.run_cell_magic("sparksql", "-q -c -e -v vce out", "SELECT 7 AS a") + assert displayed == [] + assert spark.table("vce").collect()[0].a == 7 + assert ip.user_ns["out"].collect()[0].a == 7 + + +def test_explain(ip, displayed): + out = io.StringIO() + with redirect_stdout(out): + ip.run_cell_magic("sparksql", "--explain", "SELECT 1 AS a") + assert "Physical Plan" in out.getvalue() + assert displayed == [] + + +def test_named_session(ip, spark, displayed): + ip.user_ns["other"] = spark + ip.run_cell_magic("sparksql", "-s other", "SELECT 1 AS a") + assert "1 row" in displayed[0]["text/html"] + + +def test_error_with_caret(ip): + cell = "SELECT 1;\nSELECT\n nope FROM range(1)" + with pytest.raises(sparksql.SparkSqlError) as info: + ip.run_cell_magic("sparksql", "", cell) + lines = info.value.message.splitlines() + assert lines[0].startswith("[UNRESOLVED_COLUMN") + assert "SQLSTATE: 42703" in lines[0] + assert lines[1] == ( + "(line 3; statement 2 of 2; earlier statements already ran)" + ) + assert lines[2:] == [" nope FROM range(1)", " ^^^^"] + assert isinstance(info.value.__cause__, PySparkException) + + +def test_parse_error(ip): + with pytest.raises(sparksql.SparkSqlError) as info: + ip.run_cell_magic("sparksql", "", "SELECT 1;\n\n SELEC 2") + assert info.value.message.splitlines() == [ + "[PARSE_SYNTAX_ERROR] Syntax error at or near 'SELEC'. SQLSTATE: 42601", + "(line 3; statement 2 of 2; earlier statements already ran)", + " SELEC 2", + " ^^^^^", + ] + + +def test_full_traceback(ip): + sparksql.get_config(ip).full_traceback = True + with pytest.raises(PySparkException): + ip.run_cell_magic("sparksql", "", "SELECT nope") diff --git a/tests/unit/managed_spark_magics/test_sparksql.py b/tests/unit/managed_spark_magics/test_sparksql.py new file mode 100644 index 00000000..4cec5898 --- /dev/null +++ b/tests/unit/managed_spark_magics/test_sparksql.py @@ -0,0 +1,907 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for %%sparksql that need neither a JVM nor a Spark server. + +End-to-end behaviour against real classic and Connect sessions is covered by +test_sparksql_local.py. +""" + +import unittest +from unittest import mock + +from IPython.core.error import UsageError +from IPython.core.interactiveshell import InteractiveShell +from pyspark.errors import PySparkException, QueryContextType +from traitlets import TraitError +from traitlets.config import Config, Configurable + +from google.cloud.managed_spark_magics import sparksql + + +class ParseOptionsTest(unittest.TestCase): + + def test_cell_defaults(self): + opts, rest = sparksql.parse_options("", line_form=False) + self.assertEqual(opts, sparksql.Options()) + self.assertEqual(rest, "") + + def test_cell_all_options(self): + opts, _ = sparksql.parse_options( + "-s spark -l 5 -v tmp -c -e -q --explain cost df", line_form=False + ) + self.assertEqual( + opts, + sparksql.Options( + session="spark", + limit=5, + view="tmp", + cache=True, + eager=True, + quiet=True, + explain="cost", + variable="df", + ), + ) + + def test_long_and_attached_forms(self): + opts, _ = sparksql.parse_options( + "--session=s --limit=7 --view v --cache -ce -l9", line_form=False + ) + self.assertEqual(opts.session, "s") + self.assertEqual(opts.limit, 9) + self.assertEqual(opts.view, "v") + self.assertTrue(opts.cache) + self.assertTrue(opts.eager) + + def test_explain_mode_is_optional(self): + opts, _ = sparksql.parse_options("--explain df", line_form=False) + self.assertEqual(opts.explain, "formatted") + self.assertEqual(opts.variable, "df") + opts, _ = sparksql.parse_options("--explain=EXTENDED", line_form=False) + self.assertEqual(opts.explain, "extended") + + def test_errors(self): + for line, msg in [ + ("--explain=bogus", "Unknown explain mode"), + ("-l", "needs a value"), + ("-l 0", "at least 1"), + ("-l x", "must be an integer"), + ("-x", "Unknown option"), + ("a b", "Unexpected argument"), + ("my-df", "Unknown option|Not a valid"), + ("1df", "Not a valid Python identifier"), + ("-s 1x", "Not a valid Python identifier"), + ("--view=", "needs a value"), + ("-v=", "needs a value"), + ("--session=", "needs a value"), + ("-v my.view", "Not a valid temporary view name"), + ("-v 1x", "Not a valid temporary view name"), + ]: + with self.subTest(line=line): + with self.assertRaisesRegex(UsageError, msg): + sparksql.parse_options(line, line_form=False) + + def test_usage_in_errors(self): + with self.assertRaisesRegex(UsageError, "Usage:"): + sparksql.parse_options("-x", line_form=False) + + def test_line_form(self): + opts, rest = sparksql.parse_options( + "-c SELECT -1 AS x -- c", line_form=True + ) + self.assertTrue(opts.cache) + self.assertTrue(opts.quiet) + # The SQL is everything from the first non-option token, verbatim. + self.assertEqual(rest, "SELECT -1 AS x -- c") + + def test_line_form_needs_sql(self): + with self.assertRaisesRegex(UsageError, "No SQL given"): + sparksql.parse_options("-c", line_form=True) + + def test_line_form_rejects_display_options(self): + # The line form never displays, so these would be silently ignored. + for line in ("-l 3 SELECT 1",): + with self.subTest(line=line): + with self.assertRaisesRegex(UsageError, "to the cell form"): + sparksql.parse_options(line, line_form=True) + + +class SplitStatementsTest(unittest.TestCase): + + def test_offsets_and_stripping(self): + cell = "SELECT 1;\n\n SELECT ';' AS s ;\n" + stmts = sparksql.split_statements(cell) + self.assertEqual( + [s.text for s in stmts], ["SELECT 1", "SELECT ';' AS s"] + ) + for s in stmts: + self.assertEqual(cell[s.offset : s.offset + len(s.text)], s.text) + + def test_comments_only_and_empty(self): + self.assertEqual(sparksql.split_statements(""), []) + self.assertEqual(sparksql.split_statements("-- hi\n/* x */ ;\n"), []) + + def test_comment_only_statement_is_dropped(self): + stmts = sparksql.split_statements("SELECT 1; -- trailing\n") + self.assertEqual([s.text for s in stmts], ["SELECT 1"]) + + def test_script_is_not_split(self): + cell = "-- a script\nBEGIN\n SELECT 1;\n SELECT 2;\nEND;\n" + stmts = sparksql.split_statements(cell) + self.assertEqual(len(stmts), 1) + self.assertTrue(stmts[0].text.startswith("-- a script")) + self.assertTrue(stmts[0].text.endswith("END")) + + +class ParametersTest(unittest.TestCase): + + def setUp(self): + mock.patch.object( + sparksql, "_dataframe_classes", return_value=(FakeDataFrame,) + ).start() + self.addCleanup(mock.patch.stopall) + self.df = FakeDataFrame("SELECT 1", None, {}) + + def _stmt(self, sql): + (stmt,) = sparksql.split_statements(sql) + return stmt + + def test_find_parameters(self): + stmt = self._stmt( + "SELECT :a, :b, :a, ':c', x::int, {df}, {col.name} -- :d {e}\n" + "FROM IDENTIFIER(:t) /* :f */" + ) + markers, fields = sparksql.find_parameters(stmt) + self.assertEqual(markers, ["a", "b", "t"]) + self.assertEqual(fields, ["df"]) + + def test_braces_in_strings_and_non_identifiers_are_escaped(self): + stmt = self._stmt("SELECT '{x}', \"}\", {a.b}, {0} FROM {df}") + self.assertEqual( + sparksql.escape_braces(stmt), + "SELECT '{{x}}', \"}}\", {{a.b}}, {{0}} FROM {df}", + ) + + def test_braces_found_whatever_their_token_type(self): + # sqlparse currently emits braces as Error tokens; nothing should + # depend on that, e.g. if a lexer change made them Punctuation. + import dataclasses + + from sqlparse import tokens as T + + stmt = self._stmt( + "SELECT '{x}', `{y}`, {a.b}, :n FROM {df} -- {z}\nWHERE {df2}.c" + ) + self.assertIn((T.Error, "{"), stmt.tokens) + stmt = dataclasses.replace( + stmt, + tokens=[ + (T.Punctuation if v in ("{", "}") else ttype, v) + for ttype, v in stmt.tokens + ], + ) + self.assertEqual(sparksql.find_parameters(stmt), (["n"], ["df", "df2"])) + self.assertEqual( + sparksql.escape_braces(stmt), + "SELECT '{{x}}', `{{y}}`, {{a.b}}, :n FROM {df} -- {{z}}\n" + "WHERE {df2}.c", + ) + + def test_bind(self): + stmt = self._stmt("SELECT :n, '{lit}' FROM {df};") + sql, args, kwargs = sparksql.bind_parameters( + stmt, {"n": 1, "df": self.df, "lit": "unused"} + ) + self.assertEqual(sql, "SELECT :n, '{{lit}}' FROM {df}") + self.assertEqual(args, {"n": 1}) + self.assertEqual(kwargs, {"df": self.df}) + + def test_bind_without_kwargs_keeps_text(self): + stmt = self._stmt("SELECT '{x}', :n") + sql, args, kwargs = sparksql.bind_parameters(stmt, {"n": 2}) + self.assertEqual(sql, "SELECT '{x}', :n") + self.assertEqual(args, {"n": 2}) + self.assertEqual(kwargs, {}) + + def test_bind_missing(self): + stmt = self._stmt("SELECT :a, :b FROM {c}") + with self.assertRaisesRegex(UsageError, "Not defined.*: b, c"): + sparksql.bind_parameters(stmt, {"a": 1}) + + def test_bind_rejects_non_dataframe_in_braces(self): + stmt = self._stmt("SELECT * FROM {t}") + with self.assertRaisesRegex( + UsageError, + r"\{t\}: expected a Spark DataFrame, got str\." + r" Use :t .* or IDENTIFIER\(:t\)", + ): + sparksql.bind_parameters(stmt, {"t": "my_table"}) + + # A non-Spark class also named DataFrame (e.g. pandas.DataFrame) + # includes its package name so the message isn't "got DataFrame". + FakePandasDF = type( + "DataFrame", (), {"__module__": "pandas.core.frame"} + ) + with self.assertRaisesRegex( + UsageError, r"expected a Spark DataFrame, got pandas\.DataFrame" + ): + sparksql.bind_parameters(stmt, {"t": FakePandasDF()}) + + def test_bind_rejects_dataframe_in_marker(self): + stmt = self._stmt("SELECT :df") + with self.assertRaisesRegex( + UsageError, + r":df: expected a SQL value, got a Spark DataFrame\. Use \{df\}", + ): + sparksql.bind_parameters(stmt, {"df": self.df}) + + +class ResolveSessionTest(unittest.TestCase): + + class FakeSession: + pass + + def setUp(self): + patcher = mock.patch.object( + sparksql, "_session_classes", return_value=(self.FakeSession,) + ) + patcher.start() + self.addCleanup(patcher.stop) + self.active = mock.patch.object( + sparksql, "_active_session", return_value=None + ).start() + self.addCleanup(mock.patch.stopall) + + def test_named(self): + s = self.FakeSession() + self.active.return_value = self.FakeSession() + self.assertIs(sparksql.resolve_session({"s": s}, "s"), s) + + def test_named_errors(self): + with self.assertRaisesRegex(UsageError, "no such variable"): + sparksql.resolve_session({}, "s") + with self.assertRaisesRegex(UsageError, "expected a SparkSession"): + sparksql.resolve_session({"s": 1}, "s") + + def test_active_wins_over_namespace(self): + active = self.FakeSession() + self.active.return_value = active + self.assertIs( + sparksql.resolve_session({"other": self.FakeSession()}, None), + active, + ) + + def test_only_session_in_namespace(self): + s = self.FakeSession() + self.assertIs( + sparksql.resolve_session({"s": s, "_s": self.FakeSession()}, None), + s, + ) + + def test_ambiguous_and_missing(self): + with self.assertRaisesRegex(UsageError, r"several.*\(a, b\)"): + sparksql.resolve_session( + {"b": self.FakeSession(), "a": self.FakeSession()}, None + ) + with self.assertRaisesRegex(UsageError, "No SparkSession found"): + sparksql.resolve_session({}, None) + + +class FakeContext: + + def __init__(self, start, stop, fragment=""): + self._start, self._stop, self._fragment = start, stop, fragment + + def contextType(self): + return QueryContextType.SQL + + def startIndex(self): + return self._start + + def stopIndex(self): + return self._stop + + def fragment(self): + # Empty, like a context that does not say, unless a test sets it. + return self._fragment + + +class FakeSparkError(PySparkException): + + def __init__(self, message, condition=None, sqlstate=None, span=None): + """`span` is (start, stop) or (start, stop, fragment).""" + super().__init__(message=message) + self._msg, self._cond, self._state = message, condition, sqlstate + self._ctx = [FakeContext(*span)] if span else [] + + def getMessage(self): + return self._msg + + def getCondition(self): + return self._cond + + def getSqlState(self): + return self._state + + def getQueryContext(self): + return self._ctx + + +class FormatErrorTest(unittest.TestCase): + + def test_caret_in_second_statement(self): + cell = "SELECT 1;\nSELECT\n nope FROM t" + stmts = sparksql.split_statements(cell) + exc = FakeSparkError( + "[UNRESOLVED_COLUMN] Column `nope` cannot be resolved." + " SQLSTATE: 42703\n== SQL (line 2, position 3) ==\n...", + condition="UNRESOLVED_COLUMN", + sqlstate="42703", + span=(9, 12), + ) + msg = sparksql.format_error(exc, cell, stmts[1], 1, 2, True) + self.assertEqual( + msg.splitlines(), + [ + "[UNRESOLVED_COLUMN] Column `nope` cannot be resolved." + " SQLSTATE: 42703", + "(line 3; statement 2 of 2; earlier statements already ran)", + " nope FROM t", + " ^^^^", + ], + ) + + def test_adds_condition_and_sqlstate(self): + (stmt,) = sparksql.split_statements("SELECT x") + exc = FakeSparkError("Something broke.", "SOME_CONDITION", "XX000") + msg = sparksql.format_error(exc, "SELECT x", stmt, 0, 1, True) + self.assertEqual( + msg, "[SOME_CONDITION] Something broke. SQLSTATE: XX000" + ) + + def test_no_caret_when_offsets_do_not_match(self): + (stmt,) = sparksql.split_statements("SELECT x FROM {df}") + exc = FakeSparkError("Bad.", span=(7, 7)) + msg = sparksql.format_error( + exc, "SELECT x FROM {df}", stmt, 0, 1, False + ) + self.assertEqual(msg, "Bad.") + + def test_caret_when_fragment_matches(self): + (stmt,) = sparksql.split_statements("SELECT 1/0 FROM t") + exc = FakeSparkError("Bad.", span=(7, 9, "1/0")) + msg = sparksql.format_error(exc, "SELECT 1/0 FROM t", stmt, 0, 1, True) + self.assertEqual( + msg.splitlines(), + ["Bad.", "(line 1)", " SELECT 1/0 FROM t", " ^^^"], + ) + + def test_no_caret_when_fragment_is_from_elsewhere(self): + # E.g. an error inside a view: the offsets are into the view's SQL. + (stmt,) = sparksql.split_statements("SELECT * FROM some_view") + exc = FakeSparkError("Bad.", span=(7, 9, "1/0")) + msg = sparksql.format_error( + exc, "SELECT * FROM some_view", stmt, 0, 1, True + ) + self.assertEqual(msg, "Bad.") + + def test_drops_jvm_stacktrace(self): + (stmt,) = sparksql.split_statements("SELECT x") + exc = FakeSparkError("Bad.\n\nJVM stacktrace:\norg.apache...") + self.assertEqual( + sparksql.format_error(exc, "SELECT x", stmt, 0, 1, True), "Bad." + ) + + def test_drops_connect_plan(self): + (stmt,) = sparksql.split_statements("SELECT * FROM t") + exc = FakeSparkError( + "[TABLE_OR_VIEW_NOT_FOUND] Not found.\nVerify the spelling." + " SQLSTATE: 42P01; line 1 pos 14;\n'Project [*]\n+- 'Unresolved" + ) + self.assertEqual( + sparksql.format_error(exc, "SELECT * FROM t", stmt, 0, 1, True), + "[TABLE_OR_VIEW_NOT_FOUND] Not found.\nVerify the spelling." + " SQLSTATE: 42P01", + ) + + def test_parse_error_position(self): + cell = "SELECT 1;\nSELECT\n 1 +* 2" + stmts = sparksql.split_statements(cell) + exc = FakeSparkError( + "\n[PARSE_SYNTAX_ERROR] Syntax error at or near '*'. SQLSTATE:" + " 42601 (line 2, pos 5)\n\n== SQL ==\nSELECT\n 1 +* 2\n-----^^^\n" + ) + self.assertEqual( + sparksql.format_error(exc, cell, stmts[1], 1, 2, True).splitlines(), + [ + "[PARSE_SYNTAX_ERROR] Syntax error at or near '*'." + " SQLSTATE: 42601", + "(line 3; statement 2 of 2; earlier statements already ran)", + " 1 +* 2", + " ^", + ], + ) + + def test_render_traceback(self): + err = sparksql.SparkSqlError("a\nb") + self.assertEqual(err._render_traceback_(), ["a", "b"]) + + +class RenderTest(unittest.TestCase): + + def setUp(self): + self.table = sparksql.Table( + names=["id", ""], + types=["bigint", "string"], + numeric=[True, False], + rows=[["1", "a&b"], ["22", None]], + truncated=False, + ) + + def test_html(self): + out = sparksql.render_html(self.table) + self.assertIn("<name>", out) + self.assertIn('bigint', out) + self.assertIn('22', out) + self.assertIn('NULL', out) + self.assertIn("a&b", out) + self.assertIn("2 rows", out) + self.assertNotIn("", out) + + def test_html_style_is_scoped(self): + out = sparksql.render_html(self.table) + cls = out.split('class="', 1)[1].split('"', 1)[0] + self.assertTrue(cls.startswith("sparksql-")) + self.assertIn(f".{cls} td.n", out) + + def test_truncation_footer(self): + self.table.truncated = True + self.assertIn( + "Showing the first 2 rows; more are available", + sparksql.render_html(self.table), + ) + self.table.rows = self.table.rows[:1] + self.table.truncated = False + self.assertIn("1 row<", sparksql.render_html(self.table)) + + def test_text(self): + self.table.rows.append(["3", "x" * 50 + "\ny"]) + self.assertEqual( + sparksql.render_text(self.table).splitlines(), + [ + "+--+----------------------------------------+", + "|id| |", + "+--+----------------------------------------+", + "| 1| a&b|", + "|22| NULL|", + "| 3|" + "x" * 37 + "...|", + "+--+----------------------------------------+", + "3 rows", + ], + ) + + +class FakeQueryDataFrame: + """Records the toDF/selectExpr/limit chain fetch_table() builds.""" + + def __init__(self, fields, rows): + from pyspark.sql.types import StructType + + self.schema = StructType(fields) + self._rows = rows + self.renamed = self.exprs = self.limit_n = None + + def toDF(self, *names): + self.renamed = list(names) + return self + + def selectExpr(self, *exprs): + self.exprs = list(exprs) + return self + + def limit(self, n): + self.limit_n = n + return self + + def collect(self): + return [tuple(r) for r in self._rows[: self.limit_n]] + + +class FetchTableTest(unittest.TestCase): + + def setUp(self): + from pyspark.sql.types import ( + BinaryType, + DecimalType, + LongType, + StringType, + StructField, + ) + + self.fields = [ + StructField("id", LongType()), + StructField("id", StringType()), # Duplicate names are fine. + StructField("b", BinaryType()), + StructField("d", DecimalType(10, 2)), + ] + self.rows = [[str(i), f"s{i}", "[01]", None] for i in range(5)] + + def test_formats_server_side_with_positional_names(self): + df = FakeQueryDataFrame(self.fields, self.rows) + table = sparksql.fetch_table(df, 10) + self.assertEqual(df.renamed, ["_c0", "_c1", "_c2", "_c3"]) + self.assertEqual( + df.exprs, + [ + "CAST(_c0 AS STRING) AS _c0", + "CAST(_c1 AS STRING) AS _c1", + # Binary shows as hex bytes, like df.show(). + "CAST(concat('[', regexp_replace(hex(_c2), '(..)(?=.)'," + " '$1 '), ']') AS STRING) AS _c2", + "CAST(_c3 AS STRING) AS _c3", + ], + ) + self.assertEqual(table.names, ["id", "id", "b", "d"]) + self.assertEqual( + table.types, ["bigint", "string", "binary", "decimal(10,2)"] + ) + self.assertEqual(table.numeric, [True, False, False, True]) + self.assertEqual(table.rows, self.rows) + self.assertFalse(table.truncated) + + def test_fetches_one_extra_row_to_detect_truncation(self): + df = FakeQueryDataFrame(self.fields, self.rows) + table = sparksql.fetch_table(df, 3) + self.assertEqual(df.limit_n, 4) + self.assertEqual(table.rows, self.rows[:3]) + self.assertTrue(table.truncated) + + df = FakeQueryDataFrame(self.fields, self.rows) + table = sparksql.fetch_table(df, 5) + self.assertEqual(table.rows, self.rows) + self.assertFalse(table.truncated) + + +class FakeDataFrame: + + def __init__(self, sql, args, kwargs, fields=("x",)): + self.sql, self.args, self.kwargs = sql, args, kwargs + self.schema = mock.Mock(fields=list(fields)) + self.calls = [] + + def _record(self, name, *args, **kwargs): + self.calls.append((name, args, kwargs)) + return self + + def localCheckpoint(self, **kwargs): + return self._record("localCheckpoint", **kwargs) + + def cache(self): + return self._record("cache") + + def createOrReplaceTempView(self, name): + self._record("createOrReplaceTempView", name) + + def explain(self, **kwargs): + self._record("explain", **kwargs) + + +class FakeSession: + + def __init__(self, fail_on=None): + self.queries = [] + self.fail_on = fail_on + + def sql(self, sql, args=None, **kwargs): + self.queries.append(sql) + if self.fail_on is not None and self.fail_on in sql: + raise FakeSparkError("Nope.", "BAD", "42000", span=(0, 5)) + return FakeDataFrame(sql, args, kwargs) + + +class MagicTest(unittest.TestCase): + """Drives the magic through SparkSqlMagics with a fake session.""" + + def setUp(self): + self.shell = mock.create_autospec(InteractiveShell, instance=True) + self.shell.user_ns = {} + self.shell.config = Config() + self.shell.configurables = [] + self.magics = sparksql.SparkSqlMagics(shell=self.shell) + self.session = FakeSession() + mock.patch.object( + sparksql, "_active_session", return_value=self.session + ).start() + mock.patch.object( + sparksql, "_session_classes", return_value=(FakeSession,) + ).start() + mock.patch.object( + sparksql, "_dataframe_classes", return_value=(FakeDataFrame,) + ).start() + self.display = mock.patch.object(sparksql, "_display").start() + self.addCleanup(mock.patch.stopall) + + def test_config_is_registered(self): + cfgs = [ + c + for c in self.shell.configurables + if isinstance(c, sparksql.SparkSql) + ] + self.assertEqual(len(cfgs), 1) + cfg = cfgs[0] + self.assertIs(sparksql.get_config(self.shell), cfg) + self.assertEqual(cfg.limit, 20) + self.assertFalse(cfg.full_traceback) + + def test_config_limit_must_be_positive(self): + cfg = sparksql.get_config(self.shell) + for bad in (0, -1): + with self.subTest(limit=bad): + with self.assertRaisesRegex(TraitError, "at least 1"): + cfg.limit = bad + cfg.limit = 1 + self.assertEqual(cfg.limit, 1) + + def test_get_config_matches_config_from_reloaded_module(self): + # A SparkSql created by an earlier copy of the module, as after + # importlib.reload(): same module and name, but a different class. + StaleSparkSql = type( + "SparkSql", (Configurable,), {"__module__": sparksql.__name__} + ) + stale = StaleSparkSql() + self.assertNotIsInstance(stale, sparksql.SparkSql) + self.shell.configurables[:] = [stale] + self.assertIs(sparksql.get_config(self.shell), stale) + self.assertEqual(len(self.shell.configurables), 1) + + def test_get_config_without_configurable_parent(self): + shell = mock.Mock(spec=["user_ns", "configurables"]) + shell.configurables = [] + cfg = sparksql.get_config(shell) + self.assertIsNone(cfg.parent) + self.assertEqual(shell.configurables, [cfg]) + + def test_get_config_reports_bad_shell_config(self): + ip = InteractiveShell.instance() + self.addCleanup(InteractiveShell.clear_instance) + ip.configurables[:] = [ + c for c in ip.configurables if not isinstance(c, sparksql.SparkSql) + ] + ip.config.SparkSql.limit = 0 + self.addCleanup(ip.config.pop, "SparkSql", None) + with self.assertRaises(TraitError): + sparksql.get_config(ip) + + def test_cell_runs_all_and_displays_last(self): + input_df = FakeDataFrame("SELECT 0", None, {}) + self.shell.user_ns.update(n=3, df=input_df) + result = self.magics.sparksql( + "out", "CREATE VIEW v AS SELECT 1;\nSELECT :n FROM {df}" + ) + self.assertIsNone(result) + self.assertEqual( + self.session.queries, + ["CREATE VIEW v AS SELECT 1", "SELECT :n FROM {df}"], + ) + df = self.shell.user_ns["out"] + self.assertEqual(df.args, {"n": 3}) + self.assertEqual(df.kwargs, {"df": input_df}) + self.display.assert_called_once_with(df, 20) + + def test_limit_option_and_config(self): + self.magics.sparksql("-l 5", "SELECT 1") + self.display.assert_called_with(mock.ANY, 5) + sparksql.get_config(self.shell).limit = 7 + self.magics.sparksql("", "SELECT 1") + self.display.assert_called_with(mock.ANY, 7) + + def test_view_cache_eager_explain(self): + self.magics.sparksql("-e -c -v tmp --explain d", "SELECT 1") + df = self.shell.user_ns["d"] + self.assertEqual( + df.calls, + [ + # No cache(): the checkpoint is already stored. + ("localCheckpoint", (), {"eager": True}), + ("createOrReplaceTempView", ("tmp",), {}), + ("explain", (), {"mode": "formatted"}), + ], + ) + self.display.assert_not_called() + + def test_cache_without_eager(self): + self.magics.sparksql("-c -q d", "SELECT 1") + self.assertEqual(self.shell.user_ns["d"].calls, [("cache", (), {})]) + + def test_quiet(self): + self.magics.sparksql("-q d", "SELECT 1") + self.display.assert_not_called() + self.assertIn("d", self.shell.user_ns) + + def test_line_form_returns_dataframe(self): + input_df = FakeDataFrame("SELECT 0", None, {}) + self.shell.user_ns["x"] = input_df + df = self.magics.sparksql("-c SELECT * FROM {x}") + self.assertIsInstance(df, FakeDataFrame) + self.assertEqual(df.sql, "SELECT * FROM {x}") + self.assertEqual(df.kwargs, {"x": input_df}) + self.display.assert_not_called() + + def test_missing_variable(self): + with self.assertRaisesRegex(UsageError, "Not defined.*: x"): + self.magics.sparksql("SELECT {x}") + self.assertEqual(self.session.queries, []) + + def test_missing_variable_in_later_statement_runs_nothing(self): + with self.assertRaisesRegex(UsageError, "Not defined.*: y"): + self.magics.sparksql("", "CREATE VIEW v AS SELECT 1;\nSELECT :y") + self.assertEqual(self.session.queries, []) + + def test_display_error_has_caret(self): + self.display.side_effect = FakeSparkError( + "Bad.", "BAD", "42000", span=(7, 9, "1/0") + ) + with self.assertRaises(sparksql.SparkSqlError) as cm: + self.magics.sparksql("", "SELECT 1/0") + self.assertIn(" SELECT 1/0\n ^^^", cm.exception.message) + + def test_display_error_after_dataframe_reference_has_no_caret(self): + # Spark rewrote {d} before running it, so its positions are not the + # cell's. + self.shell.user_ns["d"] = FakeDataFrame("SELECT 0", None, {}) + self.display.side_effect = FakeSparkError( + "Bad.", "BAD", "42000", span=(0, 5) + ) + with self.assertRaises(sparksql.SparkSqlError) as cm: + self.magics.sparksql("", "SELECT * FROM {d}") + self.assertEqual(cm.exception.message, "[BAD] Bad. SQLSTATE: 42000") + + def test_variable_not_bound_when_display_fails(self): + self.display.side_effect = FakeSparkError("Bad.") + with self.assertRaises(sparksql.SparkSqlError): + self.magics.sparksql("out", "SELECT 1") + self.assertNotIn("out", self.shell.user_ns) + + def test_error_in_later_statement(self): + self.session.fail_on = "boom" + with self.assertRaises(sparksql.SparkSqlError) as cm: + self.magics.sparksql("", "SELECT 1;\nboom;\nSELECT 3") + self.assertEqual(self.session.queries, ["SELECT 1", "boom"]) + msg = str(cm.exception) + self.assertIn("[BAD] Nope. SQLSTATE: 42000", msg) + self.assertIn("statement 2 of 3", msg) + self.assertIn(" boom;\n ^^^^", msg) + self.assertIsInstance(cm.exception.__cause__, FakeSparkError) + + def test_full_traceback(self): + self.session.fail_on = "boom" + sparksql.get_config(self.shell).full_traceback = True + with self.assertRaises(FakeSparkError): + self.magics.sparksql("", "boom") + + def test_non_spark_errors_propagate(self): + self.session.sql = mock.Mock(side_effect=KeyError("k")) + with self.assertRaises(KeyError): + self.magics.sparksql("", "SELECT 1") + + def test_empty_cell(self): + with self.assertRaisesRegex(UsageError, "No SQL given"): + self.magics.sparksql("", "-- nothing\n") + + +class ExtensionTest(unittest.TestCase): + + _EXTENSIONS = ( + "google.cloud.managed_spark_magics", + "google.cloud.managed_spark_magics.sparksql", + ) + + def _reset(self): + magics = self.ip.magics_manager.magics + for kind in ("line", "cell"): + magics.get(kind, {}).pop("sparksql", None) + magics.get("line", {}).pop("dpip", None) + for ext in self._EXTENSIONS: + self.ip.extension_manager.loaded.discard(ext) + + def setUp(self): + self.ip = InteractiveShell.instance() + self._reset() + + def tearDown(self): + self._reset() + InteractiveShell.clear_instance() + + def test_standalone_load_and_unload_ext(self): + magics = self.ip.magics_manager.magics + self.ip.run_line_magic( + "load_ext", "google.cloud.managed_spark_magics.sparksql" + ) + self.assertIn("sparksql", magics["line"]) + self.assertIn("sparksql", magics["cell"]) + self.assertNotIn("dpip", magics["line"]) + + self.ip.run_line_magic( + "unload_ext", "google.cloud.managed_spark_magics.sparksql" + ) + self.assertNotIn("sparksql", magics["line"]) + self.assertNotIn("sparksql", magics["cell"]) + + def test_package_unload_keeps_separately_loaded_sparksql(self): + magics = self.ip.magics_manager.magics + self.ip.run_line_magic( + "load_ext", "google.cloud.managed_spark_magics.sparksql" + ) + self.ip.run_line_magic("load_ext", "google.cloud.managed_spark_magics") + self.ip.run_line_magic( + "unload_ext", "google.cloud.managed_spark_magics" + ) + self.assertNotIn("dpip", magics["line"]) + # Still loaded in its own right, so still registered. + self.assertIn("sparksql", magics["line"]) + self.assertIn("sparksql", magics["cell"]) + + self.ip.run_line_magic( + "unload_ext", "google.cloud.managed_spark_magics.sparksql" + ) + self.assertNotIn("sparksql", magics["cell"]) + + def test_package_unload_removes_its_own_sparksql(self): + magics = self.ip.magics_manager.magics + self.ip.run_line_magic("load_ext", "google.cloud.managed_spark_magics") + self.ip.run_line_magic( + "unload_ext", "google.cloud.managed_spark_magics" + ) + self.assertNotIn("dpip", magics["line"]) + self.assertNotIn("sparksql", magics["line"]) + self.assertNotIn("sparksql", magics["cell"]) + + +_MISSING_DEPS_PROBE = """ +import sys + +for name in ("sqlparse", "pyspark"): + sys.modules[name] = None + +from IPython.core.error import UsageError +from IPython.core.interactiveshell import InteractiveShell + +ip = InteractiveShell.instance() +ip.run_line_magic("load_ext", "google.cloud.managed_spark_magics") +print("dpip" in ip.magics_manager.magics["line"]) +try: + ip.run_cell_magic("sparksql", "", "SELECT 1") +except UsageError as e: + print("interactive]" in str(e)) +""" + + +class MissingDependenciesTest(unittest.TestCase): + + def test_extension_loads_without_sqlparse_or_pyspark(self): + import subprocess + import sys + + result = subprocess.run( + [sys.executable, "-c", _MISSING_DEPS_PROBE], + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(result.stdout.split(), ["True", "True"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index 19422110..2101d787 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -221,6 +221,9 @@ def test_successful_extras_loading(self, mock_get_ipython): self.assertIn("Loaded interactive extras", output) self.assertIn("explore_dataframe()", output) self.assertIn("%dpip", output) + # A line+cell magic is listed once, under its cell name. + self.assertIn("%%sparksql", output) + self.assertEqual(output.count("sparksql"), 1) self.assertIn("ManagedSparkConnect.enable_extras = False", output) # One summary line, not one per extra. self.assertEqual(len(output.splitlines()), 1) @@ -651,8 +654,22 @@ def test_module_import_calls_init_extras(self): mock_init.assert_called_once() +_OUR_MAGICS = (("line", "dpip"), ("line", "sparksql"), ("cell", "sparksql")) + + +def _pop_our_magics(ip): + for kind, name in _OUR_MAGICS: + ip.magics_manager.magics.get(kind, {}).pop(name, None) + + class TestExtrasInteractiveShell(unittest.TestCase): + def _assert_magics(self, present): + magics = self.ip.magics_manager.magics + for kind, name in _OUR_MAGICS: + with self.subTest(kind=kind, name=name): + self.assertEqual(name in magics[kind], present) + def setUp(self): import os from IPython.core.interactiveshell import InteractiveShell @@ -666,10 +683,13 @@ def setUp(self): self.ip = InteractiveShell.instance() self.ip.user_ns.pop("explore_dataframe", None) - self.ip.magics_manager.magics.get("line", {}).pop("dpip", None) + _pop_our_magics(self.ip) self.ip.extension_manager.loaded.discard( "google.cloud.managed_spark_magics" ) + self.ip.extension_manager.loaded.discard( + "google.cloud.managed_spark_magics.sparksql" + ) if "ManagedSparkConnect" in self.ip.config: del self.ip.config["ManagedSparkConnect"] self.ip.configurables[:] = [ @@ -689,10 +709,13 @@ def tearDown(self): self._get_ipython_patcher.stop() self.ip.user_ns.pop("explore_dataframe", None) - self.ip.magics_manager.magics.get("line", {}).pop("dpip", None) + _pop_our_magics(self.ip) self.ip.extension_manager.loaded.discard( "google.cloud.managed_spark_magics" ) + self.ip.extension_manager.loaded.discard( + "google.cloud.managed_spark_magics.sparksql" + ) if "ManagedSparkConnect" in self.ip.config: del self.ip.config["ManagedSparkConnect"] self.ip.configurables[:] = [ @@ -711,14 +734,14 @@ def test_real_shell_loads_all_extras_and_toggles_via_config(self): _init_extras() self.assertIn("explore_dataframe", self.ip.user_ns) - self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=True) self.ip.run_line_magic( "config", "ManagedSparkConnect.enable_extras = False" ) self.assertNotIn("explore_dataframe", self.ip.user_ns) - self.assertNotIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=False) with mock.patch("sys.stdout", io.StringIO()): self.ip.run_line_magic( @@ -726,7 +749,7 @@ def test_real_shell_loads_all_extras_and_toggles_via_config(self): ) self.assertIn("explore_dataframe", self.ip.user_ns) - self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=True) def test_real_shell_config_magic_pre_and_post_import(self): import io @@ -740,7 +763,7 @@ def test_real_shell_config_magic_pre_and_post_import(self): _init_extras() self.assertNotIn("explore_dataframe", self.ip.user_ns) - self.assertNotIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=False) # 2. Enable via %config AFTER _init_extras() with mock.patch("sys.stdout", io.StringIO()): @@ -749,7 +772,7 @@ def test_real_shell_config_magic_pre_and_post_import(self): ) self.assertIn("explore_dataframe", self.ip.user_ns) - self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=True) # 3. Disable via %config AFTER _init_extras() self.ip.run_line_magic( @@ -757,7 +780,7 @@ def test_real_shell_config_magic_pre_and_post_import(self): ) self.assertNotIn("explore_dataframe", self.ip.user_ns) - self.assertNotIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=False) def test_real_shell_partial_undo_preserves_preloaded_extension(self): import io @@ -767,14 +790,14 @@ def test_real_shell_partial_undo_preserves_preloaded_extension(self): self.ip.extension_manager.load_extension( "google.cloud.managed_spark_magics" ) - self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=True) stdout_capture = io.StringIO() with mock.patch("sys.stdout", stdout_capture): _init_extras() self.assertIn("explore_dataframe", self.ip.user_ns) - self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=True) # %dpip was the user's doing, so the summary must not claim it. output = stdout_capture.getvalue() self.assertIn("explore_dataframe()", output) @@ -786,7 +809,103 @@ def test_real_shell_partial_undo_preserves_preloaded_extension(self): ) self.assertNotIn("explore_dataframe", self.ip.user_ns) - self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=True) + + def test_real_shell_preserves_preloaded_sparksql_extension(self): + import io + from google.cloud.managed_spark_connect._ipython import ( + _get_or_create_config, + _init_extras, + ) + + self.ip.extension_manager.load_extension( + "google.cloud.managed_spark_magics.sparksql" + ) + magics = self.ip.magics_manager.magics + line_fn = magics["line"]["sparksql"] + cell_fn = magics["cell"]["sparksql"] + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + self.assertIn("explore_dataframe", self.ip.user_ns) + self.assertIn("dpip", magics["line"]) + self.assertIs(magics["line"]["sparksql"], line_fn) + self.assertIs(magics["cell"]["sparksql"], cell_fn) + output = stdout_capture.getvalue() + self.assertIn("%dpip", output) + self.assertNotIn("sparksql", output) + + _get_or_create_config(self.ip).enable_extras = False + + self.assertNotIn("explore_dataframe", self.ip.user_ns) + self.assertNotIn("dpip", magics["line"]) + self.assertIs(magics["line"]["sparksql"], line_fn) + self.assertIs(magics["cell"]["sparksql"], cell_fn) + + def test_real_shell_preserves_other_sparksql_magic(self): + import io + from google.cloud.managed_spark_connect._ipython import ( + _get_or_create_config, + _init_extras, + ) + + def other_sparksql(line, cell=None): + pass + + magics = self.ip.magics_manager.magics + self.ip.register_magic_function( + other_sparksql, magic_kind="line_cell", magic_name="sparksql" + ) + line_fn = magics["line"]["sparksql"] + cell_fn = magics["cell"]["sparksql"] + + stdout_capture = io.StringIO() + with mock.patch("sys.stdout", stdout_capture): + _init_extras() + + # %dpip is still ours to add, but the existing %%sparksql stays put + # and is not claimed in the summary. + self.assertIn("dpip", magics["line"]) + self.assertIs(magics["line"]["sparksql"], line_fn) + self.assertIs(magics["cell"]["sparksql"], cell_fn) + output = stdout_capture.getvalue() + self.assertIn("%dpip", output) + self.assertNotIn("sparksql", output) + + _get_or_create_config(self.ip).enable_extras = False + + self.assertNotIn("dpip", magics["line"]) + self.assertIs(magics["line"]["sparksql"], line_fn) + self.assertIs(magics["cell"]["sparksql"], cell_fn) + + def test_real_shell_unload_keeps_sparksql_registered_later(self): + import io + from google.cloud.managed_spark_connect._ipython import ( + _get_or_create_config, + _init_extras, + ) + + with mock.patch("sys.stdout", io.StringIO()): + _init_extras() + + # Something else takes over %%sparksql after the extras loaded. + def other_sparksql(line, cell=None): + pass + + magics = self.ip.magics_manager.magics + self.ip.register_magic_function( + other_sparksql, magic_kind="line_cell", magic_name="sparksql" + ) + line_fn = magics["line"]["sparksql"] + cell_fn = magics["cell"]["sparksql"] + + _get_or_create_config(self.ip).enable_extras = False + + self.assertNotIn("dpip", magics["line"]) + self.assertIs(magics["line"]["sparksql"], line_fn) + self.assertIs(magics["cell"]["sparksql"], cell_fn) def test_real_shell_config_magic_overrides_env_opt_out(self): import io @@ -806,7 +925,7 @@ def test_real_shell_config_magic_overrides_env_opt_out(self): _init_extras() self.assertIn("explore_dataframe", self.ip.user_ns) - self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=True) def test_real_shell_env_opt_out_then_enable_via_config_magic(self): import io @@ -829,7 +948,7 @@ def test_real_shell_env_opt_out_then_enable_via_config_magic(self): ) self.assertIn("explore_dataframe", self.ip.user_ns) - self.assertIn("dpip", self.ip.magics_manager.magics["line"]) + self._assert_magics(present=True) _IMPORT_ORDER_PROBE = """ @@ -847,6 +966,8 @@ def test_real_shell_env_opt_out_then_enable_via_config_magic(self): print("explore_dataframe" in ip.user_ns) print("dpip" in ip.magics_manager.magics["line"]) +print("sparksql" in ip.magics_manager.magics["line"]) +print("sparksql" in ip.magics_manager.magics["cell"]) """ @@ -881,8 +1002,8 @@ def test_all_import_orders_load_all_extras(self): # which silently dropped %dpip. self.assertEqual( self._run(order), - ["True", "True"], - "expected explore_dataframe and %dpip", + ["True", "True", "True", "True"], + "expected explore_dataframe, %dpip and %%sparksql", ) @@ -891,7 +1012,7 @@ def test_all_import_orders_load_all_extras(self): # Simulate an install without the [interactive] extra. A None entry makes both # importlib.util.find_spec() and import treat the module as absent. -for name in ("IPython", "traitlets", "google.colabsqlviz"): +for name in ("IPython", "sqlparse", "traitlets", "google.colabsqlviz"): sys.modules[name] = None import google.cloud.managed_spark_connect as msc From 2ee8f82cc2c99d19a1809386dbde062c6bc2aae8 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Tue, 6 Oct 2026 12:43:22 -0700 Subject: [PATCH 3/4] docs: move bundled magics documentation out of the README The README had grown a long %%sparksql reference, while %dpip had no documentation beyond a one-line mention. Move each bundled magic's reference into its own page, docs/dpip.md and docs/sparksql.md, and leave the README with a short Bundled Magics section: one example per magic plus a link to its page. The links use absolute GitHub URLs, since the README is also the PyPI long description, where relative links don't resolve. --- README.md | 102 ++++++++------------------------------------ docs/dpip.md | 35 +++++++++++++++ docs/sparksql.md | 108 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 161 insertions(+), 84 deletions(-) create mode 100644 docs/dpip.md create mode 100644 docs/sparksql.md diff --git a/README.md b/README.md index bb5a03a6..135812cf 100644 --- a/README.md +++ b/README.md @@ -135,8 +135,7 @@ up a few interactive conveniences: - `explore_dataframe()` from [google-colabsqlviz](https://pypi.org/project/google-colabsqlviz/) is injected into your notebook globals. -- The `%dpip` line magic is loaded. -- The [`%%sparksql`](#spark-sql-magic-sparksql) magic is loaded. +- The [`%dpip` and `%%sparksql`](#bundled-magics) magics are loaded. ```python import google.cloud.managed_spark_connect # extras load here @@ -171,101 +170,36 @@ Or at runtime: An explicit IPython setting takes precedence over the environment variable. -### Spark SQL Magic (`%%sparksql`) +### Bundled Magics -`%%sparksql` runs Spark SQL against the active session, classic or Spark -Connect, and displays the result as a table. It is one of the interactive -extras, so with the `interactive` extra installed it is loaded on import; you -can also load it explicitly with `%load_ext google.cloud.managed_spark_magics`. +These are loaded with the interactive extras, or explicitly with +`%load_ext google.cloud.managed_spark_magics`. -```python -%%sparksql -SELECT * FROM your_table -``` +#### `%dpip` -The cell displays the first 20 rows by default. Change that limit per cell with -`-l N`, or globally for the notebook with: +Installs PyPI packages on the session's driver and executors, so your UDFs can +import them: ```python -%config SparkSql.limit = 100 +%dpip install spacy==3.8.4 torch ``` -Options go on the first line, optionally naming a variable to store the -resulting DataFrame in. When a variable name is provided, it receives the full, -un-truncated `DataFrame` (the `-l` display limit only affects what is shown on -screen): - -```python -%%sparksql --cache --view result_view df -SELECT * FROM your_table WHERE condition = true -``` +See [docs/dpip.md](https://github.com/GoogleCloudDataproc/managed-spark-connect-python/blob/main/docs/dpip.md) +for details. -- `-l N` / `--limit N`: Maximum rows to display (default: `SparkSql.limit`, 20). -- `-v VIEW` / `--view VIEW`: Register or replace a temporary view with this name. -- `-c` / `--cache`: Cache the resulting DataFrame (lazily). -- `-e` / `--eager`: Compute the result immediately via `localCheckpoint(eager=True)`. -- `-q` / `--quiet`: Suppress table display. -- `--explain [MODE]`: Print the query plan instead of displaying results (`simple`, - `extended`, `codegen`, `cost`, or `formatted`, the default). -- `-s NAME` / `--session NAME`: Run against the SparkSession in variable `NAME`. - Otherwise the active session is used, or else the only SparkSession in the notebook. +#### `%%sparksql` -To return the DataFrame directly for assignment or method chaining instead of -displaying a table, use the line magic form `%sparksql` (which is always quiet): +Runs Spark SQL against the active session and displays the result as a table, +with `:name` binding notebook variables as SQL values: ```python -df = %sparksql SELECT * FROM your_table +%%sparksql +SELECT * FROM IDENTIFIER(:table) WHERE id > :min_id ``` -#### Passing Python values - -Notebook variables are passed into queries using either `:name` or `{name}`: - -- **`:name` (SQL literals)**: Binds the notebook variable `name` as a typed SQL - parameter marker (via `spark.sql(args=...)`), passed directly to Spark - without string interpolation or SQL injection risk. -- **`IDENTIFIER(:name)` (Catalog & column identifiers)**: In SQL grammar, - parameter markers only substitute for literal values, not table or column - identifiers. Use Spark's built-in `IDENTIFIER()` clause to parameterize table, - view, or column names from a string variable: - - ```python - table, min_id = "my_catalog.db.events", 100 - ``` - - ```python - %%sparksql - SELECT * FROM IDENTIFIER(:table) WHERE id > :min_id - ``` - -- **`{name}` (DataFrames)**: References a Spark `DataFrame` notebook variable - to query as a relation: - - ```python - %%sparksql - SELECT country, count(*) FROM {events_df} GROUP BY country - ``` - -`:name` and `{name}` are only recognized in SQL expression positions; text -inside string literals (e.g. `'{not_a_variable}'` or `':id'`) and SQL comments -is preserved verbatim. - -#### Multiple statements and errors - -A cell can hold multiple statements separated by `;`. They run in sequence, -and only the final statement's result is displayed or stored in the target -variable. If an intermediate statement fails, earlier statements have already -executed. A cell starting with `BEGIN` is passed through whole as a Spark SQL -compound script. - -When a statement fails, `%%sparksql` prints a clean summary with Spark's error -condition, message, SQLSTATE, and a caret pointing at the offending line in the -cell, without JVM or Python wrapper tracebacks. To re-raise the original -exception with the full traceback, set: - -```python -%config SparkSql.full_traceback = True -``` +See [docs/sparksql.md](https://github.com/GoogleCloudDataproc/managed-spark-connect-python/blob/main/docs/sparksql.md) +for options, passing Python values (including `IDENTIFIER()` for table names) +and more. ## Migrating from dataproc-spark-connect diff --git a/docs/dpip.md b/docs/dpip.md new file mode 100644 index 00000000..69f3db56 --- /dev/null +++ b/docs/dpip.md @@ -0,0 +1,35 @@ +# `%dpip`: Install PyPI Packages into a Managed Spark Session + +`%dpip` installs PyPI packages, with their dependencies, on the driver and +executors of the Managed Spark session in your notebook, so that UDFs and other +code running in Spark can import them. + +```python +%dpip install spacy==3.8.4 torch +``` + +It is one of the [interactive extras](../README.md#interactive-extras), so it is +loaded when you import `google.cloud.managed_spark_connect` in a notebook. You +can also load it explicitly: + +```python +%load_ext google.cloud.managed_spark_magics +``` + +## Usage + +``` +%dpip install [ ...] +``` + +- Each package is a pip requirement: a name, optionally with a version + specifier such as `==3.8.4` (quote any specifier containing `<` or `>`). +- pip options such as `--upgrade` are not supported. +- The notebook must have exactly one `ManagedSparkSession` in a variable; + `%dpip` installs into that session. + +`%dpip` only affects the remote session. To use a package in the notebook +itself, install it locally with `%pip` as well. + +`%dpip` is equivalent to calling +`spark.addArtifacts("spacy==3.8.4", "torch", pypi=True)` on the session. diff --git a/docs/sparksql.md b/docs/sparksql.md new file mode 100644 index 00000000..bbcf77d3 --- /dev/null +++ b/docs/sparksql.md @@ -0,0 +1,108 @@ +# `%%sparksql`: Spark SQL Magic + +`%%sparksql` runs Spark SQL against the active session, classic or Spark +Connect, and displays the result as a table. + +```python +%%sparksql +SELECT * FROM your_table +``` + +It is one of the [interactive extras](../README.md#interactive-extras), so it is +loaded when you import `google.cloud.managed_spark_connect` in a notebook. You +can also load it explicitly: + +```python +%load_ext google.cloud.managed_spark_magics +``` + +## Options + +The cell displays the first 20 rows by default. Change that limit per cell with +`-l N`, or globally for the notebook with: + +```python +%config SparkSql.limit = 100 +``` + +Options go on the first line, optionally naming a variable to store the +resulting DataFrame in. When a variable name is provided, it receives the full, +un-truncated `DataFrame` (the `-l` display limit only affects what is shown on +screen): + +```python +%%sparksql --cache --view result_view df +SELECT * FROM your_table WHERE condition = true +``` + +- `-l N` / `--limit N`: Maximum rows to display (default: `SparkSql.limit`, 20). +- `-v VIEW` / `--view VIEW`: Register or replace a temporary view with this name. +- `-c` / `--cache`: Cache the resulting DataFrame (lazily). Redundant with `-e`, + whose checkpoint is already stored. +- `-e` / `--eager`: Compute the result immediately via `localCheckpoint(eager=True)`. +- `-q` / `--quiet`: Suppress table display. +- `--explain [MODE]`: Print the query plan instead of displaying results (`simple`, + `extended`, `codegen`, `cost`, or `formatted`, the default). +- `-s NAME` / `--session NAME`: Run against the SparkSession in variable `NAME`. + Otherwise the active session is used, or else the only SparkSession in the notebook. + +### Line form (`%sparksql`) + +As a line magic, `%sparksql` takes options followed by the SQL query on the +same line. It is always `--quiet` (it does not display the table, so `-d` and +`-l` are rejected) and returns the resulting DataFrame directly instead of +taking a variable name argument: + +```python +df = %sparksql SELECT * FROM your_table +``` + +## Passing Python values + +Notebook variables are passed into queries using either `:name` or `{name}`: + +- **`:name` (SQL literals)**: Binds the notebook variable `name` as a typed SQL + parameter marker (via `spark.sql(args=...)`), passed directly to Spark + without string interpolation or SQL injection risk. +- **`IDENTIFIER(:name)` (Catalog & column identifiers)**: In SQL grammar, + parameter markers only substitute for literal values, not table or column + identifiers. Use Spark's built-in `IDENTIFIER()` clause to parameterize table, + view, or column names from a string variable: + + ```python + table, min_id = "my_catalog.db.events", 100 + ``` + + ```python + %%sparksql + SELECT * FROM IDENTIFIER(:table) WHERE id > :min_id + ``` + +- **`{name}` (DataFrames)**: References a Spark `DataFrame` notebook variable + to query as a relation: + + ```python + %%sparksql + SELECT country, count(*) FROM {events_df} GROUP BY country + ``` + +`:name` and `{name}` are only recognized in SQL expression positions; text +inside string literals (e.g. `'{not_a_variable}'` or `':id'`) and SQL comments +is preserved verbatim. + +## Multiple statements and errors + +A cell can hold multiple statements separated by `;`. They run in sequence, +and only the final statement's result is displayed or stored in the target +variable. If an intermediate statement fails, earlier statements have already +executed. A cell starting with `BEGIN` is passed through whole as a Spark SQL +compound script. + +When a statement fails, `%%sparksql` prints a clean summary with Spark's error +condition, message, SQLSTATE, and a caret pointing at the offending line in the +cell, without JVM or Python wrapper tracebacks. To re-raise the original +exception with the full traceback, set: + +```python +%config SparkSql.full_traceback = True +``` From 5e06bad62bec9525005fe5bfbcc5aa422644edd4 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Tue, 6 Oct 2026 13:28:07 -0700 Subject: [PATCH 4/4] feat: render %%sparksql results with colabsqlviz by default Upgrade google-colabsqlviz to >=0.4.0, which supports chart_type='table' on explore_dataframe(), and route %%sparksql's result display through it by default so results open as an interactive table with pagination and chart switching. - Add -d / --display (colabsqlviz, viz, or html) and %config SparkSql.display (defaulting to 'colabsqlviz') to choose between colabsqlviz and the built-in HTML/text table renderer. - When colabsqlviz is the default and google-colabsqlviz is not installed, fall back to the HTML renderer; if -d colabsqlviz / -d viz was explicitly passed on the cell, raise a UsageError pointing to the [interactive] extra. - Default SparkSql.limit to None so colabsqlviz paginates all rows unless capped via -l N or %config SparkSql.limit = N (which passes df.limit(N)), while the HTML renderer still defaults to 20 rows. --- docs/sparksql.md | 18 +- google/cloud/managed_spark_magics/sparksql.py | 109 +++++++++++-- requirements-dev.txt | 2 +- setup.py | 2 +- tests/local/test_sparksql_local.py | 38 ++++- .../managed_spark_magics/test_sparksql.py | 154 +++++++++++++++++- tests/unit/test_init.py | 2 +- 7 files changed, 289 insertions(+), 36 deletions(-) diff --git a/docs/sparksql.md b/docs/sparksql.md index bbcf77d3..34d11075 100644 --- a/docs/sparksql.md +++ b/docs/sparksql.md @@ -18,8 +18,16 @@ can also load it explicitly: ## Options -The cell displays the first 20 rows by default. Change that limit per cell with -`-l N`, or globally for the notebook with: +By default, results are rendered as an interactive table with `colabsqlviz`, +which paginates through all rows and lets you switch to charts. To render a +plain HTML table instead (showing the first 20 rows by default), pass `-d html` +per cell or configure it for the whole notebook: + +```python +%config SparkSql.display = "html" +``` + +To cap the number of rows displayed, pass `-l N` per cell or set: ```python %config SparkSql.limit = 100 @@ -35,7 +43,11 @@ screen): SELECT * FROM your_table WHERE condition = true ``` -- `-l N` / `--limit N`: Maximum rows to display (default: `SparkSql.limit`, 20). +- `-d MODE` / `--display MODE`: How to display the result (`colabsqlviz` / + `viz`, the default, or `html`; configurable via `SparkSql.display`). If + `google-colabsqlviz` is not installed, `%%sparksql` falls back to `html`. +- `-l N` / `--limit N`: Maximum rows to display (default: `SparkSql.limit`; + when unset, `colabsqlviz` paginates all rows and `html` shows 20). - `-v VIEW` / `--view VIEW`: Register or replace a temporary view with this name. - `-c` / `--cache`: Cache the resulting DataFrame (lazily). Redundant with `-e`, whose checkpoint is already stored. diff --git a/google/cloud/managed_spark_magics/sparksql.py b/google/cloud/managed_spark_magics/sparksql.py index 080b7a0d..fd4c1f00 100644 --- a/google/cloud/managed_spark_magics/sparksql.py +++ b/google/cloud/managed_spark_magics/sparksql.py @@ -46,17 +46,19 @@ magics_class, no_var_expand, ) -from traitlets import Bool, Int, TraitError, validate +from traitlets import Bool, CaselessStrEnum, Int, TraitError, validate from traitlets.config import Configurable USAGE = """\ -%%sparksql [-s SESSION] [-l LIMIT] [-v VIEW] [-c] [-e] [-q] [--explain [MODE]] [variable] +%%sparksql [-s SESSION] [-d DISPLAY] [-l LIMIT] [-v VIEW] [-c] [-e] [-q] [--explain [MODE]] [variable] df = %sparksql [-s SESSION] [-v VIEW] [-c] [-e] [--explain [MODE]] """ EXPLAIN_MODES = ("simple", "extended", "codegen", "cost", "formatted") +DISPLAY_MODES = ("colabsqlviz", "viz", "html") +DEFAULT_HTML_LIMIT = 20 _IDENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") _PLACEHOLDER_RE = re.compile(r":([A-Za-z_][A-Za-z0-9_]*)\Z") @@ -73,8 +75,23 @@ class SparkSql(Configurable): third-party sparksql-magic package keep working. """ + display = CaselessStrEnum( + DISPLAY_MODES, + default_value="colabsqlviz", + help=( + "How %%sparksql displays results: 'colabsqlviz' (or 'viz') for an" + " interactive table, or 'html' for a plain HTML table (-d" + " overrides)." + ), + ).tag(config=True) limit = Int( - 20, help="Default number of rows %%sparksql displays (-l overrides)." + None, + allow_none=True, + help=( + "Maximum number of rows %%sparksql displays (-l overrides). By" + " default, colabsqlviz paginates all rows and html shows" + f" {DEFAULT_HTML_LIMIT}." + ), ).tag(config=True) full_traceback = Bool( False, @@ -84,12 +101,17 @@ class SparkSql(Configurable): ), ).tag(config=True) + @validate("display") + def _validate_display(self, proposal: Mapping[str, Any]) -> str: + val = proposal["value"] + return "colabsqlviz" if val == "viz" else val + @validate("limit") def _validate_limit(self, proposal: Mapping[str, Any]) -> Optional[int]: # Same rule as -l, which parse_options enforces. val = proposal["value"] if val is not None and val < 1: - raise TraitError("SparkSql.limit must be at least 1") + raise TraitError("SparkSql.limit must be at least 1, or None") return val @@ -127,6 +149,7 @@ def get_config(shell: Any) -> SparkSql: @dataclasses.dataclass class Options: session: Optional[str] = None + display: Optional[str] = None limit: Optional[int] = None view: Optional[str] = None cache: bool = False @@ -139,6 +162,8 @@ class Options: _VALUE_OPTS = { "-s": "session", "--session": "session", + "-d": "display", + "--display": "display", "-l": "limit", "--limit": "limit", "-v": "view", @@ -232,7 +257,7 @@ def take_value(flag: str) -> str: elif ( not flag.startswith("--") and len(flag) > 2 - and flag[1] in "slv" + and flag[1] in "sdlv" and not eq ): # Attached short value, e.g. -l50. @@ -249,6 +274,14 @@ def take_value(flag: str) -> str: raise _usage_error(f"Unknown option: {tok}") i += 1 + if opts.display is not None: + mode = opts.display.lower() + if mode not in DISPLAY_MODES: + raise _usage_error( + f"Unknown display mode {opts.display!r}; expected one of" + f" {', '.join(DISPLAY_MODES)}" + ) + opts.display = "colabsqlviz" if mode == "viz" else mode if opts.limit is not None: try: limit = int(opts.limit) @@ -268,10 +301,12 @@ def take_value(flag: str) -> str: # Spark error after the query has already run. raise _usage_error(f"Not a valid temporary view name: {opts.view!r}") if line_form: - if opts.limit is not None: - # The line form never displays anything, so this would be + if opts.display is not None or opts.limit is not None: + # The line form never displays anything, so these would be # silently ignored. - raise _usage_error("-l only applies to the cell form (%%sparksql)") + raise _usage_error( + "-d and -l only apply to the cell form (%%sparksql)" + ) if not rest.strip(): raise _usage_error("No SQL given") opts.quiet = True @@ -908,7 +943,14 @@ def run(shell: Any, line: str, cell: Optional[str]) -> Any: if opts.explain is not None: df.explain(mode=opts.explain) elif not opts.quiet: - _display(df, opts.limit or config.limit) + limit = opts.limit if opts.limit is not None else config.limit + display_mode = opts.display or config.display + _display( + df, + limit=limit, + display_mode=display_mode, + explicit_display=opts.display is not None, + ) # Only once everything has succeeded, so that a failure does not # leave the variable bound to a DataFrame that does not work. if opts.variable is not None: @@ -956,16 +998,45 @@ def _reraise( raise SparkSqlError(message) from exc -def _display(df: Any, limit: int) -> None: +def _import_explore_dataframe() -> Any: + try: + from google.colabsqlviz.explore_dataframe import explore_dataframe + + return explore_dataframe + except ImportError: + return None + + +def _display( + df: Any, + limit: Optional[int] = None, + display_mode: str = "colabsqlviz", + explicit_display: bool = False, +) -> None: from IPython.display import display if not df.schema.fields: # DDL and most DML: there is nothing to show but that it worked. display({"text/plain": "OK"}, raw=True) return - # TODO: Route through colabsqlviz's explore_dataframe() once it can be - # asked to open as a table rather than a chart. - table = fetch_table(df, limit) + if display_mode in ("colabsqlviz", "viz"): + explore_dataframe = _import_explore_dataframe() + if explore_dataframe is not None: + if limit is not None: + df = df.limit(limit) + # Relies on explore_dataframe() only returning the widget, which + # it does while colabsqlviz's anywidget_enabled experiment is on + # (the default in 0.4.0). With it off, explore_dataframe() also + # displays the widget itself, and the table would render twice. + display(explore_dataframe(df, chart_type="table")) + return + if explicit_display: + raise UsageError( + "%%sparksql -d colabsqlviz needs the google-colabsqlviz" + " package. Install it with:\n" + " pip install 'google-cloud-spark-connect[interactive]'" + ) + table = fetch_table(df, limit if limit is not None else DEFAULT_HTML_LIMIT) display( {"text/html": render_html(table), "text/plain": render_text(table)}, raw=True, @@ -995,8 +1066,8 @@ def sparksql(self, line: str, cell: Optional[str] = None) -> Any: Cell form: displays the result of the last statement. - %%sparksql [-s SESSION] [-l LIMIT] [-v VIEW] [-c] [-e] [-q] - [--explain [MODE]] [variable] + %%sparksql [-s SESSION] [-d DISPLAY] [-l LIMIT] [-v VIEW] [-c] [-e] + [-q] [--explain [MODE]] [variable] SELECT * FROM t WHERE x = :x Line form: always --quiet; returns the DataFrame directly. @@ -1009,8 +1080,14 @@ def sparksql(self, line: str, cell: Optional[str] = None) -> Any: (cell form only; -l only limits displayed rows). -s, --session N Use the SparkSession in variable N. Default: the active session, else the only one in the notebook. + -d, --display M How to display results: colabsqlviz (or viz) for an + interactive table, or html for a plain HTML table + (default: %config SparkSql.display, initially + colabsqlviz; falls back to html if + google-colabsqlviz is not installed). -l, --limit N Maximum rows to display (default: %config - SparkSql.limit, initially 20). + SparkSql.limit; when unset, colabsqlviz paginates + all rows and html shows 20). -v, --view V Also register the result as temp view V. -c, --cache Cache the result (lazily). Redundant with -e. -e, --eager Materialize the result now, with diff --git a/requirements-dev.txt b/requirements-dev.txt index c4ad48ba..f3c15fcf 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,6 +1,6 @@ google-api-core>=2.19 google-cloud-dataproc>=5.18 -google-colabsqlviz>=0.3.0 +google-colabsqlviz>=0.4.0 ipython~=9.1 ipywidgets>=8.0.0 packaging>=20.0 diff --git a/setup.py b/setup.py index 0c235823..ed02808a 100644 --- a/setup.py +++ b/setup.py @@ -40,7 +40,7 @@ # Everything here is optional at runtime: managed_spark_connect._ipython # only initializes the extras whose dependencies are importable. "interactive": [ - "google-colabsqlviz>=0.3.0", + "google-colabsqlviz>=0.4.0", # Imported directly by managed_spark_connect._ipython and # managed_spark_magics; previously these only arrived transitively # via google-colabsqlviz. diff --git a/tests/local/test_sparksql_local.py b/tests/local/test_sparksql_local.py index 12e79f89..b52315d9 100644 --- a/tests/local/test_sparksql_local.py +++ b/tests/local/test_sparksql_local.py @@ -105,7 +105,8 @@ def ip(spark): shell.run_line_magic( "load_ext", "google.cloud.managed_spark_magics.sparksql" ) - sparksql.get_config(shell).limit = 20 + sparksql.get_config(shell).display = "colabsqlviz" + sparksql.get_config(shell).limit = None sparksql.get_config(shell).full_traceback = False shell.user_ns["spark"] = spark yield shell @@ -126,10 +127,36 @@ def displayed(): yield calls -def test_select_renders_table(ip, displayed): +def test_select_renders_colabsqlviz_table(ip, displayed): + # colabsqlviz has no public way to inspect a rendered viz, so this reads + # its internals (_df, init_data_b64, lsapi_pb2) and may need updating + # along with the google-colabsqlviz version. + import base64 + from google.colabsqlviz import interactive_viz + from google.colabsqlviz.lsapi import lsapi_pb2 + + # Default limit is unlimited (more than DEFAULT_HTML_LIMIT = 20 rows). ip.run_cell_magic( "sparksql", "", + "SELECT id, CAST(id AS STRING) AS s FROM range(25) ORDER BY id", + ) + (viz,) = displayed + assert isinstance(viz, interactive_viz.InteractiveViz) + init_data = lsapi_pb2.ColabInitData.FromString( + base64.b64decode(viz.init_data_b64) + ) + assert init_data.chart_hints.chart_type == "TABLE" + assert len(viz._df) == 25 + + ip.run_cell_magic("sparksql", "-l 2", "SELECT * FROM range(25)") + assert len(displayed[1]._df) == 2 + + +def test_select_renders_html_table(ip, displayed): + ip.run_cell_magic( + "sparksql", + "-d html", "SELECT id, CAST(id AS STRING) AS s FROM range(3) ORDER BY id", ) (bundle,) = displayed @@ -141,7 +168,8 @@ def test_select_renders_table(ip, displayed): assert bundle["text/plain"].splitlines()[1] == "|id|s|" -def test_limit_and_truncation(ip, displayed): +def test_html_limit_and_truncation(ip, displayed): + sparksql.get_config(ip).display = "html" ip.run_cell_magic("sparksql", "-l 2", "SELECT * FROM range(5)") assert "Showing the first 2 rows" in displayed[0]["text/html"] sparksql.get_config(ip).limit = 5 @@ -152,7 +180,7 @@ def test_limit_and_truncation(ip, displayed): def test_values_are_formatted_by_spark(ip, displayed): ip.run_cell_magic( "sparksql", - "", + "-d html", "SELECT CAST('bin' AS BINARY) AS b, CAST(NULL AS INT) AS n," " array(1, NULL) AS a, map('k', 1.5D) AS m, DATE'2026-01-02' AS d", ) @@ -241,7 +269,7 @@ def test_explain(ip, displayed): def test_named_session(ip, spark, displayed): ip.user_ns["other"] = spark - ip.run_cell_magic("sparksql", "-s other", "SELECT 1 AS a") + ip.run_cell_magic("sparksql", "-s other -d html", "SELECT 1 AS a") assert "1 row" in displayed[0]["text/html"] diff --git a/tests/unit/managed_spark_magics/test_sparksql.py b/tests/unit/managed_spark_magics/test_sparksql.py index 4cec5898..6183c67f 100644 --- a/tests/unit/managed_spark_magics/test_sparksql.py +++ b/tests/unit/managed_spark_magics/test_sparksql.py @@ -39,12 +39,14 @@ def test_cell_defaults(self): def test_cell_all_options(self): opts, _ = sparksql.parse_options( - "-s spark -l 5 -v tmp -c -e -q --explain cost df", line_form=False + "-s spark -d html -l 5 -v tmp -c -e -q --explain cost df", + line_form=False, ) self.assertEqual( opts, sparksql.Options( session="spark", + display="html", limit=5, view="tmp", cache=True, @@ -57,14 +59,19 @@ def test_cell_all_options(self): def test_long_and_attached_forms(self): opts, _ = sparksql.parse_options( - "--session=s --limit=7 --view v --cache -ce -l9", line_form=False + "--session=s --display=VIZ --limit=7 --view v --cache -ce -l9", + line_form=False, ) self.assertEqual(opts.session, "s") + self.assertEqual(opts.display, "colabsqlviz") self.assertEqual(opts.limit, 9) self.assertEqual(opts.view, "v") self.assertTrue(opts.cache) self.assertTrue(opts.eager) + opts, _ = sparksql.parse_options("-dhtml", line_form=False) + self.assertEqual(opts.display, "html") + def test_explain_mode_is_optional(self): opts, _ = sparksql.parse_options("--explain df", line_form=False) self.assertEqual(opts.explain, "formatted") @@ -75,6 +82,8 @@ def test_explain_mode_is_optional(self): def test_errors(self): for line, msg in [ ("--explain=bogus", "Unknown explain mode"), + ("--display=bogus", "Unknown display mode"), + ("-d", "needs a value"), ("-l", "needs a value"), ("-l 0", "at least 1"), ("-l x", "must be an integer"), @@ -112,9 +121,11 @@ def test_line_form_needs_sql(self): def test_line_form_rejects_display_options(self): # The line form never displays, so these would be silently ignored. - for line in ("-l 3 SELECT 1",): + for line in ("-l 3 SELECT 1", "-d html SELECT 1"): with self.subTest(line=line): - with self.assertRaisesRegex(UsageError, "to the cell form"): + with self.assertRaisesRegex( + UsageError, "only apply to the cell form" + ): sparksql.parse_options(line, line_form=True) @@ -585,6 +596,9 @@ def _record(self, name, *args, **kwargs): self.calls.append((name, args, kwargs)) return self + def limit(self, n): + return self._record("limit", n) + def localCheckpoint(self, **kwargs): return self._record("localCheckpoint", **kwargs) @@ -598,6 +612,83 @@ def explain(self, **kwargs): self._record("explain", **kwargs) +class DisplayTest(unittest.TestCase): + + def setUp(self): + self.df = FakeDataFrame("SELECT 1", None, {}) + self.ipython_display = mock.patch("IPython.display.display").start() + self.explore = mock.Mock(return_value="VIZ") + self.import_explore = mock.patch.object( + sparksql, "_import_explore_dataframe", return_value=self.explore + ).start() + self.table = sparksql.Table( + names=["x"], + types=["bigint"], + numeric=[True], + rows=[["1"]], + truncated=False, + ) + self.fetch_table = mock.patch.object( + sparksql, "fetch_table", return_value=self.table + ).start() + self.addCleanup(mock.patch.stopall) + + def test_empty_schema_displays_ok(self): + empty = FakeDataFrame("CREATE VIEW v AS SELECT 1", None, {}, fields=()) + sparksql._display(empty) + self.ipython_display.assert_called_once_with( + {"text/plain": "OK"}, raw=True + ) + self.explore.assert_not_called() + self.fetch_table.assert_not_called() + + def test_colabsqlviz_default_is_unlimited(self): + sparksql._display(self.df) + self.assertEqual(self.df.calls, []) + self.explore.assert_called_once_with(self.df, chart_type="table") + self.ipython_display.assert_called_once_with("VIZ") + self.fetch_table.assert_not_called() + + def test_colabsqlviz_with_limit(self): + sparksql._display(self.df, limit=5) + self.assertEqual(self.df.calls, [("limit", (5,), {})]) + self.explore.assert_called_once_with(self.df, chart_type="table") + self.ipython_display.assert_called_once_with("VIZ") + + def test_colabsqlviz_falls_back_to_html_when_not_installed(self): + self.import_explore.return_value = None + sparksql._display(self.df) + self.fetch_table.assert_called_once_with( + self.df, sparksql.DEFAULT_HTML_LIMIT + ) + self.ipython_display.assert_called_once() + bundle = self.ipython_display.call_args[0][0] + self.assertIn("text/html", bundle) + self.assertIn("text/plain", bundle) + + def test_colabsqlviz_explicit_raises_when_not_installed(self): + self.import_explore.return_value = None + with self.assertRaisesRegex( + UsageError, + r"pip install 'google-cloud-spark-connect\[interactive\]'", + ): + sparksql._display( + self.df, display_mode="colabsqlviz", explicit_display=True + ) + self.fetch_table.assert_not_called() + + def test_html_mode_default_and_custom_limit(self): + sparksql._display(self.df, display_mode="html") + self.fetch_table.assert_called_once_with( + self.df, sparksql.DEFAULT_HTML_LIMIT + ) + self.explore.assert_not_called() + + self.fetch_table.reset_mock() + sparksql._display(self.df, limit=7, display_mode="html") + self.fetch_table.assert_called_once_with(self.df, 7) + + class FakeSession: def __init__(self, fail_on=None): @@ -642,8 +733,13 @@ def test_config_is_registered(self): self.assertEqual(len(cfgs), 1) cfg = cfgs[0] self.assertIs(sparksql.get_config(self.shell), cfg) - self.assertEqual(cfg.limit, 20) + self.assertEqual(cfg.display, "colabsqlviz") + self.assertIsNone(cfg.limit) self.assertFalse(cfg.full_traceback) + cfg.display = "VIZ" + self.assertEqual(cfg.display, "colabsqlviz") + cfg.display = "HTML" + self.assertEqual(cfg.display, "html") def test_config_limit_must_be_positive(self): cfg = sparksql.get_config(self.shell) @@ -652,7 +748,8 @@ def test_config_limit_must_be_positive(self): with self.assertRaisesRegex(TraitError, "at least 1"): cfg.limit = bad cfg.limit = 1 - self.assertEqual(cfg.limit, 1) + cfg.limit = None + self.assertIsNone(cfg.limit) def test_get_config_matches_config_from_reloaded_module(self): # A SparkSql created by an earlier copy of the module, as after @@ -698,14 +795,53 @@ def test_cell_runs_all_and_displays_last(self): df = self.shell.user_ns["out"] self.assertEqual(df.args, {"n": 3}) self.assertEqual(df.kwargs, {"df": input_df}) - self.display.assert_called_once_with(df, 20) + self.display.assert_called_once_with( + df, + limit=None, + display_mode="colabsqlviz", + explicit_display=False, + ) def test_limit_option_and_config(self): self.magics.sparksql("-l 5", "SELECT 1") - self.display.assert_called_with(mock.ANY, 5) + self.display.assert_called_with( + mock.ANY, + limit=5, + display_mode="colabsqlviz", + explicit_display=False, + ) sparksql.get_config(self.shell).limit = 7 self.magics.sparksql("", "SELECT 1") - self.display.assert_called_with(mock.ANY, 7) + self.display.assert_called_with( + mock.ANY, + limit=7, + display_mode="colabsqlviz", + explicit_display=False, + ) + + def test_display_option_and_config(self): + self.magics.sparksql("-d html", "SELECT 1") + self.display.assert_called_with( + mock.ANY, + limit=None, + display_mode="html", + explicit_display=True, + ) + sparksql.get_config(self.shell).display = "html" + self.magics.sparksql("", "SELECT 1") + self.display.assert_called_with( + mock.ANY, + limit=None, + display_mode="html", + explicit_display=False, + ) + self.magics.sparksql("-d viz", "SELECT 1") + self.display.assert_called_with( + mock.ANY, + limit=None, + display_mode="colabsqlviz", + explicit_display=True, + ) def test_view_cache_eager_explain(self): self.magics.sparksql("-e -c -v tmp --explain d", "SELECT 1") diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index 2101d787..cf164adc 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -150,7 +150,7 @@ def test_colabsqlviz_available(self): self.assertTrue(callable(display)) self.assertTrue(callable(explore_dataframe)) version = importlib.metadata.version("google-colabsqlviz") - self.assertGreaterEqual(Version(version), Version("0.3.0")) + self.assertGreaterEqual(Version(version), Version("0.4.0")) def test_managed_spark_magics_available(self): import google.cloud.managed_spark_magics as magics_ext