Use DataprocCommandRunner to run "pip install"

Author: tim-u

google/cloud/dataproc_magics/magics.py

@@ -25,6 +25,12 @@
 @magics_class
 class DataprocMagics(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,
@@ -45,9 +51,7 @@ def dpip(self, line):
             print(f"Installing packages: {packages}")
             output = self._run_command(packages, session)
 
-            failure_match = re.search(
-                "Pip install failed with non-zero exit code", output
-            )
+            failure_match = re.search(self.PIP_INSTALL_FAILURE_MSG, output)
             if failure_match:
                 raise RuntimeError(output)
 
@@ -89,10 +93,8 @@ def _check_preconditions(self, args):
 
     def _run_command(self, packages, session):
         command = pb2.Command()
-        command.execute_external_command.runner = (
-            "org.apache.spark.sql.artifact.DataprocCommandRunner"
-        )
-        command.execute_external_command.command = "PipInstallPackages"
+        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
@@ -104,9 +106,7 @@ def _run_command(self, packages, session):
 
             # decode the Arrow stream and return the output
             table = pa.ipc.RecordBatchStreamReader(binary_data).read_all()
-            return "\n".join(
-                str(log_line) for log_line in table.column(0).to_pylist()
-            )
+            return "\n".join(table.column(0).to_pylist())
         except (KeyError, AttributeError) as e:
             raise RuntimeError(
                 "Unexpected response structure: missing binary data."

tests/unit/dataproc_magics/test_magics.py

@@ -33,6 +33,14 @@ def setUp(self):
         self.shell.config = Config()
         self.magics = DataprocMagics(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."
@@ -82,13 +90,9 @@ def test_dpip_install_packages_success(self):
         properties = mock.Mock()
 
         # Create a pyarrow table and serialize it
-        schema = pa.schema([pa.field("output", pa.string())])
-        data = [["Collecting pandas", "Successfully installed pandas"]]
-        table = pa.Table.from_arrays(data, schema=schema)
-        sink = pa.BufferOutputStream()
-        with pa.ipc.RecordBatchStreamWriter(sink, table.schema) as writer:
-            writer.write_table(table)
-        binary_data = sink.getvalue()
+        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
@@ -109,7 +113,8 @@ def test_dpip_install_packages_success(self):
         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, "PipInstallPackages"
+            call_args.execute_external_command.command,
+            DataprocMagics.PIP_INSTALL_COMMAND,
         )
         self.assertEqual(
             call_args.execute_external_command.options["0"], "pandas"
@@ -130,18 +135,12 @@ def test_dpip_install_failure(self):
 
         # Create a mock for the properties object with failure message
         properties = mock.Mock()
-        schema = pa.schema([pa.field("output", pa.string())])
-        data = [
+        binary_data = self._create_mock_arrow_binary(
             [
-                "Pip install failed with non-zero exit code",
+                DataprocMagics.PIP_INSTALL_FAILURE_MSG,
                 "ERROR: some pip error",
             ]
-        ]
-        table = pa.Table.from_arrays(data, schema=schema)
-        sink = pa.BufferOutputStream()
-        with pa.ipc.RecordBatchStreamWriter(sink, table.schema) as writer:
-            writer.write_table(table)
-        binary_data = sink.getvalue()
+        )
 
         properties.sql_command_result.local_relation.data = binary_data
         mock_session.client.execute_command.return_value = (