Skip to content

Use DataprocCommandRunner to run "pip install" #187

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dborowitz merged 4 commits into from
Oct 23, 2026
Merged

Use DataprocCommandRunner to run "pip install" #187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d7e4fc0
Use DataprocCommandRunner to run "pip install"
tim-u Apr 8, 2026
65fd706
Use DataprocCommandRunner to run "pip install"
tim-u May 7, 2026
2729aa6
Use DataprocCommandRunner to run "pip install"
tim-u May 7, 2026
9a829ec
Merge remote-tracking branch 'origin/main' into dp-pip-with-output
tim-u Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 73 additions & 30 deletions google/cloud/managed_spark_magics/magics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,20 @@

"""Managed Spark magic implementations."""

import re
import shlex
from IPython.core.magic import (Magics, magics_class, line_magic)


@magics_class
class ManagedSparkMagics(Magics):

PIP_INSTALL_FAILURE_MSG = "Pip install failed with non-zero exit code"
DATAPROC_COMMAND_RUNNER = (
"org.apache.spark.sql.artifact.DataprocCommandRunner"
)
PIP_INSTALL_COMMAND = "PipInstallPackages"

def __init__(
self,
shell,
Expand All @@ -34,47 +41,83 @@ def dpip(self, line):
Custom magic to install pip packages as Spark Connect artifacts.
Usage: %dpip install pandas numpy
"""
try:
args = shlex.split(line)
packages, session = self._check_preconditions(args)

print(f"Installing packages: {packages}")
output = self._run_command(packages, session)

failure_match = re.search(self.PIP_INSTALL_FAILURE_MSG, output)
if failure_match:
raise RuntimeError(output)

print(output)
print("Finished installing packages.")
except Exception as e:
raise RuntimeError(f"Failed to install packages: {e}") from e

def _check_preconditions(self, args):
# 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)
if not args or args[0] != "install":
raise RuntimeError("Usage: %dpip install <package1> <package2> ...")

if not args or args[0] != "install":
raise RuntimeError(
"Usage: %dpip install <package1> <package2> ..."
)
packages = args[1:] # remove `install`

packages = args[1:] # remove `install`
if not packages:
raise RuntimeError("Error: No packages specified.")

if not packages:
raise RuntimeError("Error: No packages specified.")
if any(pkg.startswith("-") for pkg in packages):
raise RuntimeError("Error: Flags are not currently supported.")

if any(pkg.startswith("-") for pkg in packages):
raise RuntimeError("Error: Flags are not currently supported.")
sessions = [
(key, value)
for key, value in self.shell.user_ns.items()
if isinstance(value, ManagedSparkSession)
]

sessions = [
(key, value)
for key, value in self.shell.user_ns.items()
if isinstance(value, ManagedSparkSession)
]
if not sessions:
raise RuntimeError(
"Error: No active Managed Spark Session found. Please create one first."
)
if len(sessions) > 1:
raise RuntimeError(
"Error: Found more than one active Managed Spark Sessions."
)

if not sessions:
raise RuntimeError(
"Error: No active Managed Spark Session found. Please create one first."
)
if len(sessions) > 1:
raise RuntimeError(
"Error: Found more than one active Managed Spark Sessions."
)
((name, session),) = sessions
print(f"Active session found: {name}")
return packages, session

((name, session),) = sessions
print(f"Active session found: {name}")
print(f"Installing packages: {packages}")
session.addArtifacts(*packages, pypi=True)
def _run_command(self, packages, session):
# Imported here rather than at module scope, like the package's other
# pyspark imports: PySpark (and with it pyarrow) is an optional extra,
# and this extension must still load without it.
import pyarrow as pa
import pyspark.sql.connect.proto as pb2

print("Finished installing packages.")
command = pb2.Command()
command.execute_external_command.runner = self.DATAPROC_COMMAND_RUNNER
command.execute_external_command.command = self.PIP_INSTALL_COMMAND

for index, package in enumerate(packages):
command.execute_external_command.options[str(index)] = package

_, properties, _ = session.client.execute_command(command)

try:
binary_data = properties["sql_command_result"].local_relation.data

# decode the Arrow stream and return the output
table = pa.ipc.RecordBatchStreamReader(binary_data).read_all()
return "\n".join(table.column(0).to_pylist())
except (KeyError, AttributeError) as e:
raise RuntimeError(
"Unexpected response structure: missing binary data."
) from e
except Exception as e:
raise RuntimeError(f"Failed to install packages: {e}") from e
raise RuntimeError(f"Error decoding Arrow data: {e}") from e
88 changes: 79 additions & 9 deletions tests/unit/managed_spark_magics/test_magics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
from contextlib import redirect_stdout
from unittest import mock

import pyarrow as pa
import pyspark.sql.connect.proto as pb2
from google.cloud.managed_spark_connect import ManagedSparkSession
from google.cloud.managed_spark_magics import ManagedSparkMagics
from IPython.core.interactiveshell import InteractiveShell
Expand All @@ -31,6 +33,14 @@ def setUp(self):
self.shell.config = Config()
self.magics = ManagedSparkMagics(shell=self.shell)

def _create_mock_arrow_binary(self, lines: list[str]) -> bytes:
schema = pa.schema([pa.field("output", pa.string())])
table = pa.Table.from_arrays([lines], schema=schema)
sink = pa.BufferOutputStream()
with pa.ipc.RecordBatchStreamWriter(sink, table.schema) as writer:
writer.write_table(table)
return sink.getvalue()

def test_dpip_with_flags(self):
with self.assertRaisesRegex(
RuntimeError, "Error: Flags are not currently supported."
Expand Down Expand Up @@ -74,29 +84,89 @@ def test_dpip_no_packages_specified(self):

def test_dpip_install_packages_success(self):
mock_session = mock.Mock(spec=ManagedSparkSession)
mock_session.client = mock.Mock()

# Create a mock for the properties object
properties = mock.Mock()

# Create a pyarrow table and serialize it
binary_data = self._create_mock_arrow_binary(
["Collecting pandas", "Successfully installed pandas"]
)

# Set up the mock response structure
properties.sql_command_result.local_relation.data = binary_data
mock_session.client.execute_command.return_value = (
None,
{"sql_command_result": properties.sql_command_result},
None,
)

self.shell.user_ns["spark"] = mock_session

f = io.StringIO()
with redirect_stdout(f):
self.magics.dpip("install pandas numpy")

mock_session.addArtifacts.assert_called_once_with(
"pandas", "numpy", pypi=True
# Check that execute_command was called
mock_session.client.execute_command.assert_called_once()
call_args = mock_session.client.execute_command.call_args[0][0]
self.assertIsInstance(call_args, pb2.Command)
self.assertEqual(
call_args.execute_external_command.command,
ManagedSparkMagics.PIP_INSTALL_COMMAND,
)
self.assertEqual(
call_args.execute_external_command.options["0"], "pandas"
)
self.assertEqual(
call_args.execute_external_command.options["1"], "numpy"
)
self.assertEqual(mock_session.addArtifacts.call_count, 1)
self.assertIn("Finished installing packages.", f.getvalue())

def test_dpip_add_artifacts_fails(self):
output = f.getvalue()
self.assertIn("Installing packages: ['pandas', 'numpy']", output)
self.assertIn("Collecting pandas", output)
self.assertIn("Successfully installed pandas", output)
self.assertIn("Finished installing packages.", output)

def test_dpip_install_failure(self):
mock_session = mock.Mock(spec=ManagedSparkSession)
mock_session.addArtifacts.side_effect = Exception("Failed")
mock_session.client = mock.Mock()

# Create a mock for the properties object with failure message
properties = mock.Mock()
binary_data = self._create_mock_arrow_binary(
[
ManagedSparkMagics.PIP_INSTALL_FAILURE_MSG,
"ERROR: some pip error",
]
)

properties.sql_command_result.local_relation.data = binary_data
mock_session.client.execute_command.return_value = (
None,
{"sql_command_result": properties.sql_command_result},
None,
)

self.shell.user_ns["spark"] = mock_session

with self.assertRaisesRegex(
RuntimeError, "Failed to install packages: Failed"
RuntimeError, "Failed to install packages: Pip install failed"
):
self.magics.dpip("install pandas")
self.magics.dpip("install non-existent-package")

def test_dpip_unexpected_response(self):
mock_session = mock.Mock(spec=ManagedSparkSession)
mock_session.client = mock.Mock()
# Return response without 'sql_command_result'
mock_session.client.execute_command.return_value = (None, {}, None)
self.shell.user_ns["spark"] = mock_session

mock_session.addArtifacts.assert_called_once_with("pandas", pypi=True)
with self.assertRaisesRegex(
RuntimeError, "Unexpected response structure: missing binary data"
):
self.magics.dpip("install pandas")


if __name__ == "__main__":
Expand Down
Loading