diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index abaa39f..d298beb 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 c9a8a5e..06a1cae 100644 --- a/DEVELOPING.md +++ b/DEVELOPING.md @@ -39,28 +39,25 @@ env \ pytest --tb=auto -v ``` -## Testing with Magic Support +## Testing the Interactive Extras -To run tests with magic functionality, install the required dependencies manually: +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`. -```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 without Magic Support - -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. - 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 c111792..135812c 100644 --- a/README.md +++ b/README.md @@ -121,53 +121,86 @@ 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. -### Using Spark SQL Magic Commands (Jupyter Notebooks) +### Interactive Extras -The package supports the [sparksql-magic](https://github.com/cryeo/sparksql-magic) library for executing Spark SQL queries directly in Jupyter notebooks. +For notebooks and other interactive use, install the `interactive` extra: -**Installation**: To use magic commands, install the required dependencies manually: -```bash -pip install google-cloud-spark-connect -pip install IPython sparksql-magic +```sh +pip install 'google-cloud-spark-connect[interactive]' ``` -1. Load the magic extension: - ```python - %load_ext sparksql_magic - ``` +When you import the package inside an IPython kernel, it then automatically sets +up a few interactive conveniences: -2. Configure default settings (optional): - ```python - %config SparkSql.limit=20 - ``` +- `explore_dataframe()` from + [google-colabsqlviz](https://pypi.org/project/google-colabsqlviz/) is injected + into your notebook globals. +- The [`%dpip` and `%%sparksql`](#bundled-magics) magics are loaded. -3. Execute SQL queries: - ```python - %%sparksql - SELECT * FROM your_table - ``` +```python +import google.cloud.managed_spark_connect # extras load here +``` -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 - ``` +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. -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 +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. -See [sparksql-magic](https://github.com/cryeo/sparksql-magic) for more examples. +#### Opting out -**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 +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. + +### Bundled Magics + +These are loaded with the interactive extras, or explicitly with +`%load_ext google.cloud.managed_spark_magics`. + +#### `%dpip` + +Installs PyPI packages on the session's driver and executors, so your UDFs can +import them: + +```python +%dpip install spacy==3.8.4 torch +``` + +See [docs/dpip.md](https://github.com/GoogleCloudDataproc/managed-spark-connect-python/blob/main/docs/dpip.md) +for details. + +#### `%%sparksql` + +Runs Spark SQL against the active session and displays the result as a table, +with `:name` binding notebook variables as SQL values: + +```python +%%sparksql +SELECT * FROM IDENTIFIER(:table) WHERE id > :min_id +``` + +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 The `dataproc-spark-connect` package has been renamed to `google-cloud-spark-connect`. This is a breaking change with no compatibility shims — you need to update your code in the following places when you switch to the new package. diff --git a/docs/dpip.md b/docs/dpip.md new file mode 100644 index 0000000..69f3db5 --- /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 0000000..34d1107 --- /dev/null +++ b/docs/sparksql.md @@ -0,0 +1,120 @@ +# `%%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 + +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 +``` + +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 +``` + +- `-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. +- `-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 +``` diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index 2356758..06082a2 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 0000000..da7bee8 --- /dev/null +++ b/google/cloud/managed_spark_connect/_ipython.py @@ -0,0 +1,475 @@ +# 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, Sequence, Set, Tuple + +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"), ("line", "sparksql"), ("cell", "sparksql")), + ), +) + + +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, %%sparksql)." + ), + ).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) + # 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 + + +# 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 _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, magics: Sequence[Tuple[str, 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 + + # 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: + """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 _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" + 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, magics in _EXTRAS_EXTENSIONS: + try: + if ext_name in state.loaded_extensions or ( + _is_extension_loaded(ip, ext_name, magics) + ): + continue + if not _is_module_available(ext_name): + logger.debug( + "Not loading IPython extension %s: not installed", + 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) + state.added_magics[ext_name] = added + loaded.extend(_magics_display_names(added)) + 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) + loaded = getattr(ext_mgr, "loaded", None) + + 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) + except Exception: + logger.debug( + "Failed to unload IPython extension %s", + ext_name, + exc_info=True, + ) + try: + 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 magics for extension %s", + 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 79632f5..4a19cba 100644 --- a/google/cloud/managed_spark_magics/__init__.py +++ b/google/cloud/managed_spark_magics/__init__.py @@ -12,8 +12,34 @@ # 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): + 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): + 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/magics.py b/google/cloud/managed_spark_magics/magics.py index 54363ae..38ca3c7 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/google/cloud/managed_spark_magics/sparksql.py b/google/cloud/managed_spark_magics/sparksql.py new file mode 100644 index 0000000..fd4c1f0 --- /dev/null +++ b/google/cloud/managed_spark_magics/sparksql.py @@ -0,0 +1,1132 @@ +# 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, CaselessStrEnum, Int, TraitError, validate +from traitlets.config import Configurable + +USAGE = """\ +%%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") + +# 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. + """ + + 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( + 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, + help=( + "Show Spark's full exception and traceback for failed queries" + " instead of a short summary." + ), + ).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, or None") + 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 + display: 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", + "-d": "display", + "--display": "display", + "-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 "sdlv" + 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.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) + 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.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( + "-d and -l only apply 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: + 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: + 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 _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 + 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, + ) + + +# --------------------------------------------------------------------------- +# 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] [-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. + + 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. + -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; 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 + 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 5cf7026..f3c15fc 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,11 +1,13 @@ google-api-core>=2.19 google-cloud-dataproc>=5.18 +google-colabsqlviz>=0.4.0 ipython~=9.1 ipywidgets>=8.0.0 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 89637d5..ed02808 100644 --- a/setup.py +++ b/setup.py @@ -36,4 +36,18 @@ "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.4.0", + # Imported directly by managed_spark_connect._ipython and + # 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 c921713..b2cf110 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 0000000..8b13789 --- /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 0000000..b52315d --- /dev/null +++ b/tests/local/test_sparksql_local.py @@ -0,0 +1,304 @@ +# 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).display = "colabsqlviz" + sparksql.get_config(shell).limit = None + 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_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 + 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_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 + 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", + "-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", + ) + 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 -d html", "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 0000000..6183c67 --- /dev/null +++ b/tests/unit/managed_spark_magics/test_sparksql.py @@ -0,0 +1,1043 @@ +# 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 -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, + eager=True, + quiet=True, + explain="cost", + variable="df", + ), + ) + + def test_long_and_attached_forms(self): + opts, _ = sparksql.parse_options( + "--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") + 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"), + ("--display=bogus", "Unknown display mode"), + ("-d", "needs a value"), + ("-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", "-d html SELECT 1"): + with self.subTest(line=line): + with self.assertRaisesRegex( + UsageError, "only apply 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 limit(self, n): + return self._record("limit", n) + + 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 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): + 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.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) + for bad in (0, -1): + with self.subTest(limit=bad): + with self.assertRaisesRegex(TraitError, "at least 1"): + cfg.limit = bad + 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 + # 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, + 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, + 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, + 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") + 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 794f6fa..cf164ad 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -135,5 +135,917 @@ 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.4.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) + # 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) + + @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() + + +_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 + 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) + _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[:] = [ + 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) + _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[:] = [ + 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._assert_magics(present=True) + + self.ip.run_line_magic( + "config", "ManagedSparkConnect.enable_extras = False" + ) + + self.assertNotIn("explore_dataframe", self.ip.user_ns) + self._assert_magics(present=False) + + 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._assert_magics(present=True) + + 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._assert_magics(present=False) + + # 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._assert_magics(present=True) + + # 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._assert_magics(present=False) + + 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._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._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) + 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._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 + 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._assert_magics(present=True) + + 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._assert_magics(present=True) + + +_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"]) +print("sparksql" in ip.magics_manager.magics["line"]) +print("sparksql" in ip.magics_manager.magics["cell"]) +""" + + +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", "True", "True"], + "expected explore_dataframe, %dpip and %%sparksql", + ) + + +_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", "sqlparse", "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()