feat: Automatically load interactive extras on import with opt-out controls

.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
+      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 }}

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

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.

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 <package> [<package> ...]
+```
+
+- 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.

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
+```

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()