feat(providers/spark): add post_submit_commands to SparkSubmitHook for sidecar lifecycle management

providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py

@@ -21,6 +21,7 @@
 import contextlib
 import os
 import re
+import shlex
 import shutil
 import subprocess
 import tempfile
@@ -94,6 +95,10 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
                         (will overwrite any deployment mode defined in the connection's extra JSON)
     :param use_krb5ccache: if True, configure spark to use ticket cache instead of relying
         on keytab for Kerberos login
+    :param post_submit_commands: Optional list of shell commands to run after the Spark
+        job finishes (on both success and on_kill). Useful for cleaning up sidecars such
+        as Istio (e.g. ``["curl -X POST localhost:15020/quitquitquit"]``). Each command
+        is executed via the shell; failures produce a warning but do not fail the task.
     """
 
     conn_name_attr = "conn_id"
@@ -189,6 +194,7 @@ def __init__(
         deploy_mode: str | None = None,
         *,
         use_krb5ccache: bool = False,
+        post_submit_commands: list[str] | None = None,
     ) -> None:
         super().__init__()
         self._conf = conf or {}
@@ -237,6 +243,7 @@ def __init__(
         self._driver_status: str | None = None
         self._spark_exit_code: int | None = None
         self._env: dict[str, Any] | None = None
+        self._post_submit_commands: list[str] = list(post_submit_commands) if post_submit_commands else []
 
     def _resolve_should_track_driver_status(self) -> bool:
         """
@@ -545,9 +552,41 @@ def _resolve_kerberos_principal(self, principal: str | None) -> str:
             func = kerberos.get_kerberos_principal
         except AttributeError:
             # Fallback for older versions of Airflow
-            func = kerberos.get_kerberos_principle  # type: ignore[attr-defined]
+            func = getattr(kerberos, "get_kerberos_principle")
         return func(principal)
 
+    def _run_post_submit_commands(self) -> None:
+        """
+        Run any post-submit shell commands configured on this hook.
+
+        Called after the Spark job finishes (success or on_kill). Typical use case
+        is killing sidecars like Istio that don't shut down automatically.
+        Failures are logged as warnings and never raise.
+        """
+        for cmd in self._post_submit_commands:
+            self.log.debug("Running post-submit command: %s", cmd)
+            try:
+                result = subprocess.run(
+                    shlex.split(cmd),
+                    shell=False,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT,
+                    text=True,
+                    check=False,
+                    timeout=30,
+                )
+                self.log.debug("Post-submit command output:\n%s", result.stdout[:2000])
+                if result.returncode != 0:
+                    self.log.warning(
+                        "Post-submit command exited with non-zero code %d: %s",
+                        result.returncode,
+                        cmd,
+                    )
+            except subprocess.TimeoutExpired:
+                self.log.warning("Post-submit command timed out (30s): %s", cmd)
+            except Exception as exc:
+                self.log.warning("Post-submit command raised an exception: %s. Error: %s", cmd, exc)
+
     def submit(self, application: str = "", **kwargs: Any) -> None:
         """
         Remote Popen to execute the spark-submit job.
@@ -576,35 +615,38 @@ def submit(self, application: str = "", **kwargs: Any) -> None:
 
         # Check spark-submit return code. In Kubernetes mode, also check the value
         # of exit code in the log, as it may differ.
-        if returncode or (self._is_kubernetes and self._spark_exit_code != 0):
-            if self._is_kubernetes:
+        try:
+            if returncode or (self._is_kubernetes and self._spark_exit_code != 0):
+                if self._is_kubernetes:
+                    raise AirflowException(
+                        f"Cannot execute: {self._mask_cmd(spark_submit_cmd)}. Error code is: {returncode}. "
+                        f"Kubernetes spark exit code is: {self._spark_exit_code}"
+                    )
                 raise AirflowException(
-                    f"Cannot execute: {self._mask_cmd(spark_submit_cmd)}. Error code is: {returncode}. "
-                    f"Kubernetes spark exit code is: {self._spark_exit_code}"
+                    f"Cannot execute: {self._mask_cmd(spark_submit_cmd)}. Error code is: {returncode}."
                 )
-            raise AirflowException(
-                f"Cannot execute: {self._mask_cmd(spark_submit_cmd)}. Error code is: {returncode}."
-            )
 
-        self.log.debug("Should track driver: %s", self._should_track_driver_status)
+            self.log.debug("Should track driver: %s", self._should_track_driver_status)
 
-        # We want the Airflow job to wait until the Spark driver is finished
-        if self._should_track_driver_status:
-            if self._driver_id is None:
-                raise AirflowException(
-                    "No driver id is known: something went wrong when executing the spark submit command"
-                )
+            # We want the Airflow job to wait until the Spark driver is finished
+            if self._should_track_driver_status:
+                if self._driver_id is None:
+                    raise AirflowException(
+                        "No driver id is known: something went wrong when executing the spark submit command"
+                    )
 
-            # We start with the SUBMITTED status as initial status
-            self._driver_status = "SUBMITTED"
+                # We start with the SUBMITTED status as initial status
+                self._driver_status = "SUBMITTED"
 
-            # Start tracking the driver status (blocking function)
-            self._start_driver_status_tracking()
+                # Start tracking the driver status (blocking function)
+                self._start_driver_status_tracking()
 
-            if self._driver_status != "FINISHED":
-                raise AirflowException(
-                    f"ERROR : Driver {self._driver_id} badly exited with status {self._driver_status}"
-                )
+                if self._driver_status != "FINISHED":
+                    raise AirflowException(
+                        f"ERROR : Driver {self._driver_id} badly exited with status {self._driver_status}"
+                    )
+        finally:
+            self._run_post_submit_commands()
 
     def _process_spark_submit_log(self, itr: Iterator[Any]) -> None:
         """
@@ -827,3 +869,5 @@ def on_kill(self) -> None:
 
                 except kube_client.ApiException:
                     self.log.exception("Exception when attempting to kill Spark on K8s")
+
+        self._run_post_submit_commands()

providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py

@@ -84,6 +84,8 @@ class SparkSubmitOperator(BaseOperator):
                         (will overwrite any deployment mode defined in the connection's extra JSON)
     :param use_krb5ccache: if True, configure spark to use ticket cache instead of relying
                            on keytab for Kerberos login
+    :param post_submit_commands: Optional list of shell commands to run after the Spark job finishes.
+        Useful for cleaning up sidecars such as Istio. Failures produce a warning but do not fail the task.
     """
 
     template_fields: Sequence[str] = (
@@ -101,6 +103,7 @@ class SparkSubmitOperator(BaseOperator):
         "name",
         "application_args",
         "env_vars",
+        "post_submit_commands",
         "properties_file",
     )
 
@@ -137,6 +140,7 @@ def __init__(
         yarn_queue: str | None = None,
         deploy_mode: str | None = None,
         use_krb5ccache: bool = False,
+        post_submit_commands: list[str] | None = None,
         openlineage_inject_parent_job_info: bool = conf.getboolean(
             "openlineage", "spark_inject_parent_job_info", fallback=False
         ),
@@ -175,8 +179,11 @@ def __init__(
         self._yarn_queue = yarn_queue
         self._deploy_mode = deploy_mode
         self._hook: SparkSubmitHook | None = None
+        self.post_submit_commands = post_submit_commands
+        self._post_submit_commands = list(post_submit_commands) if post_submit_commands else []
         self._conn_id = conn_id
         self._use_krb5ccache = use_krb5ccache
+
         self._openlineage_inject_parent_job_info = openlineage_inject_parent_job_info
         self._openlineage_inject_transport_info = openlineage_inject_transport_info
 
@@ -229,4 +236,5 @@ def _get_hook(self) -> SparkSubmitHook:
             yarn_queue=self._yarn_queue,
             deploy_mode=self._deploy_mode,
             use_krb5ccache=self._use_krb5ccache,
+            post_submit_commands=self.post_submit_commands,
         )

providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py

@@ -1267,3 +1267,72 @@ def test_create_keytab_path_from_base64_keytab_with_existing_keytab(
         mock_open.assert_called_with(Path(f"resolved_path/airflow_keytab-{principal}"), "rb")
         _mock_open.read.assert_called_once()
         assert not _mock_open.write.called, "Keytab file should not be written"
+
+    @patch("airflow.providers.apache.spark.hooks.spark_submit.subprocess.run")
+    def test_run_post_submit_commands_success(self, mock_run):
+        """Test that post_submit_commands are run with shell=False and shlex.split."""
+        import subprocess
+        from unittest.mock import MagicMock
+
+        mock_result = MagicMock(spec=subprocess.CompletedProcess)
+        mock_result.returncode = 0
+        mock_result.stdout = "sidecar terminated"
+        mock_run.return_value = mock_result
+
+        hook = SparkSubmitHook(
+            conn_id="",
+            post_submit_commands=["curl -X POST localhost:15020/quitquitquit"],
+        )
+        hook._run_post_submit_commands()
+
+        mock_run.assert_called_once()
+        call_args = mock_run.call_args
+        assert (
+            call_args.kwargs.get("shell") is False
+            or call_args.args[1:2] == (False,)
+            or not call_args.kwargs.get("shell", True)
+        )
+        # Verify shlex.split was applied — cmd was split into a list
+        assert isinstance(call_args.args[0], list)
+        assert call_args.args[0] == ["curl", "-X", "POST", "localhost:15020/quitquitquit"]
+
+    @patch("airflow.providers.apache.spark.hooks.spark_submit.subprocess.run")
+    def test_run_post_submit_commands_nonzero_exit_warns(self, mock_run):
+        """Test that a non-zero exit code logs a warning but does not raise."""
+        import subprocess
+        from unittest.mock import MagicMock
+
+        mock_result = MagicMock(spec=subprocess.CompletedProcess)
+        mock_result.returncode = 1
+        mock_result.stdout = "error output"
+        mock_run.return_value = mock_result
+
+        hook = SparkSubmitHook(conn_id="", post_submit_commands=["false"])
+        # Should not raise
+        hook._run_post_submit_commands()
+        mock_run.assert_called_once()
+
+    @patch("airflow.providers.apache.spark.hooks.spark_submit.subprocess.run")
+    def test_run_post_submit_commands_timeout_warns(self, mock_run):
+        """Test that a timeout logs a warning but does not raise."""
+        import subprocess
+
+        mock_run.side_effect = subprocess.TimeoutExpired(cmd="sleep 100", timeout=30)
+
+        hook = SparkSubmitHook(conn_id="", post_submit_commands=["sleep 100"])
+        # Should not raise
+        hook._run_post_submit_commands()
+
+    def test_post_submit_commands_list_is_copied(self):
+        """Test that the commands list is copied, not stored by reference."""
+        original = ["cmd1", "cmd2"]
+        hook = SparkSubmitHook(conn_id="", post_submit_commands=original)
+        original.append("cmd3")
+        assert hook._post_submit_commands == ["cmd1", "cmd2"], (
+            "Hook should not be affected by mutations to the original list"
+        )
+
+    def test_post_submit_commands_none_gives_empty_list(self):
+        """Test that None post_submit_commands results in an empty list."""
+        hook = SparkSubmitHook(conn_id="")
+        assert hook._post_submit_commands == []