Expose SparkSubmitHook fields used by submit backends

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

@@ -818,7 +818,7 @@ def submit(self, application: str = "", **kwargs: Any) -> str | None:
                     f"{self._submit_log_tail}"
                 )
         finally:
-            # K8s-API tracking defers post-submit commands to _poll_k8s_driver_via_api's finally
+            # K8s-API tracking defers post-submit commands to poll_k8s_driver_via_api's finally
             # block so they run once after the driver reaches a terminal state. Spark cluster-mode
             # driver tracking defers them to poll_until_complete for the same reason. All other
             # modes run them here, immediately after spark-submit exits.
@@ -1061,7 +1061,7 @@ def _query_yarn_application_status(self, application_id: str) -> tuple[str, str,
                 f"{application_id}: {resp.text[:200]}"
             ) from exc
 
-    def _kill_yarn_application(self, application_id: str) -> None:
+    def kill_yarn_application(self, application_id: str) -> None:
         """PUT ``/ws/v1/cluster/apps/{id}/state`` to kill the application (best-effort)."""
         try:
             url = f"{self._get_yarn_rm_base_url()}/ws/v1/cluster/apps/{application_id}/state"
@@ -1175,7 +1175,7 @@ def _start_driver_status_tracking(self) -> None:
                         f"returncode = {returncode}"
                     )
 
-    def _poll_k8s_driver_via_api(self) -> str | None:
+    def poll_k8s_driver_via_api(self) -> str | None:
         """
         Poll the K8s driver pod phase until it reaches a terminal state.
 
@@ -1385,10 +1385,28 @@ def on_kill(self) -> None:
         # state because `yarn_track_via_rm_api=True` deliberately terminates
         # `_submit_sp` right after submission to free the JVM.
         if self._yarn_application_id and self._yarn_track_via_rm_api:
-            self._kill_yarn_application(self._yarn_application_id)
+            self.kill_yarn_application(self._yarn_application_id)
 
         self._run_post_submit_commands()
 
+    @property
+    def conf(self) -> dict[str, Any]:
+        """Return the live Spark conf dict (not a copy); callers may mutate keys in place."""
+        return self._conf
+
+    @property
+    def kubernetes_driver_pod(self) -> str | None:
+        return self._kubernetes_driver_pod
+
+    @kubernetes_driver_pod.setter
+    def kubernetes_driver_pod(self, value: str | None) -> None:
+        self._kubernetes_driver_pod = value
+
+    @property
+    def yarn_application_id(self) -> str | None:
+        """YARN application id captured from spark-submit logs. Get-only."""
+        return self._yarn_application_id
+
     def query_yarn_application_status(self, application_id: str) -> str:
         """
         Return a normalized single string status for the ResumableJobMixin interface.

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

@@ -96,9 +96,9 @@ class _KubernetesSparkSubmitBackend(_SparkSubmitDeploymentBackend):
     """Logic for tracking Spark driver pods in Kubernetes."""
 
     def submit_job(self, context: Context) -> str | None:
-        self.hook._conf[_K8S_WAIT_APP_COMPLETION_CONF] = "false"
+        self.hook.conf[_K8S_WAIT_APP_COMPLETION_CONF] = "false"
         self.hook.submit(self.operator.application)
-        pod_name = self.hook._kubernetes_driver_pod
+        pod_name = self.hook.kubernetes_driver_pod
         namespace = self.hook._connection["namespace"]
         if not pod_name:
             raise RuntimeError("spark-submit did not capture a K8s driver pod name")
@@ -133,8 +133,8 @@ def is_job_succeeded(self, status: str) -> bool:
     def poll_until_complete(self, external_id: str, context: Context) -> None:
         if external_id is not None:
             _, pod_name = self.operator._parse_k8s_external_id(external_id)
-            self.hook._kubernetes_driver_pod = pod_name
-        terminal_phase = self.hook._poll_k8s_driver_via_api()
+            self.hook.kubernetes_driver_pod = pod_name
+        terminal_phase = self.hook.poll_k8s_driver_via_api()
         # Cache only when the pod actually reached Succeeded, the 404/vanished path
         # returns None for cases like: pod deleted by on_kill or garbage collected after failure)
         # and must not be cached, otherwise a retry would see "Succeeded" and skip resubmission.
@@ -150,15 +150,15 @@ class _YarnSparkSubmitBackend(_SparkSubmitDeploymentBackend):
     """Logic for tracking Spark applications in YARN cluster mode."""
 
     def submit_job(self, context: Context) -> str | None:
-        if self.hook._conf.get("spark.yarn.submit.waitAppCompletion", "").strip().lower() == "true":
+        if self.hook.conf.get("spark.yarn.submit.waitAppCompletion", "").strip().lower() == "true":
             raise ValueError(
                 "spark.yarn.submit.waitAppCompletion=true cannot be set for cluster mode as it conflicts "
                 "with the need to exit spark-submit immediately to persist the application ID for tracking. "
                 "Either remove the explicit conf or set durable=False."
             )
-        self.hook._conf["spark.yarn.submit.waitAppCompletion"] = "false"
+        self.hook.conf["spark.yarn.submit.waitAppCompletion"] = "false"
         self.hook.submit(self.operator.application)
-        app_id = self.hook._yarn_application_id
+        app_id = self.hook.yarn_application_id
         if not app_id:
             raise RuntimeError("spark-submit did not produce a YARN application ID")
         self.operator.log.info("YARN application submitted: %s", app_id)
@@ -181,10 +181,10 @@ def poll_until_complete(self, external_id: str, context: Context) -> None:
             self.hook._run_post_submit_commands()
 
     def on_kill(self) -> None:
-        if self.hook._yarn_application_id:
+        if self.hook.yarn_application_id:
             # spark-submit has already exited (waitAppCompletion=false), so the hook's
             # CLI-based kill has nothing to terminate. Kill the YARN app via REST API instead.
-            self.hook._kill_yarn_application(self.hook._yarn_application_id)
+            self.hook.kill_yarn_application(self.hook.yarn_application_id)
         else:
             self.hook.on_kill()
 

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

@@ -1632,7 +1632,7 @@ def test_poll_k8s_driver_succeeds(self, mock_get_client):
         mock_client.read_namespaced_pod.side_effect = [running_pod, succeeded_pod]
 
         with patch.object(hook, "_run_post_submit_commands"):
-            hook._poll_k8s_driver_via_api()
+            hook.poll_k8s_driver_via_api()
 
         assert mock_client.delete_namespaced_pod.call_args.args[:2] == ("spark-app-abc-driver", "mynamespace")
 
@@ -1647,7 +1647,7 @@ def test_poll_k8s_driver_raises_on_failed(self, mock_get_client):
         mock_client.read_namespaced_pod.return_value = failed_pod
 
         with pytest.raises(RuntimeError, match="phase=Failed"):
-            hook._poll_k8s_driver_via_api()
+            hook.poll_k8s_driver_via_api()
 
     @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
     def test_poll_k8s_driver_raises_after_consecutive_unknown(self, mock_get_client):
@@ -1659,7 +1659,7 @@ def test_poll_k8s_driver_raises_after_consecutive_unknown(self, mock_get_client)
         mock_client.read_namespaced_pod.return_value = V1Pod(status=V1PodStatus(phase="Unknown"))
 
         with patch("time.sleep"), pytest.raises(RuntimeError, match="Unknown phase"):
-            hook._poll_k8s_driver_via_api()
+            hook.poll_k8s_driver_via_api()
 
         # assert that it was polled minimum 3 times to confirm the Unknown status before raising
         assert mock_client.read_namespaced_pod.call_count == 3
@@ -1677,13 +1677,13 @@ def test_poll_k8s_driver_tolerates_transient_api_errors(self, mock_get_client, _
         mock_client.read_namespaced_pod.side_effect = [api_error, api_error, succeeded_pod]
 
         with patch.object(hook, "_run_post_submit_commands"):
-            hook._poll_k8s_driver_via_api()
+            hook.poll_k8s_driver_via_api()
 
         assert mock_client.read_namespaced_pod.call_count == 3
 
     @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
     def test_post_submit_commands_run_exactly_once_on_k8s_path(self, mock_get_client):
-        """_run_post_submit_commands must fire exactly once: in _poll_k8s_driver_via_api finally."""
+        """_run_post_submit_commands must fire exactly once: in poll_k8s_driver_via_api finally."""
         hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True)
         hook._kubernetes_driver_pod = "spark-app-abc-driver"
         hook._kubernetes_application_id = "spark-abc"
@@ -1692,7 +1692,7 @@ def test_post_submit_commands_run_exactly_once_on_k8s_path(self, mock_get_client
         mock_client.read_namespaced_pod.return_value = V1Pod(status=V1PodStatus(phase="Succeeded"))
 
         with patch.object(hook, "_run_post_submit_commands") as mock_cmd:
-            hook._poll_k8s_driver_via_api()
+            hook.poll_k8s_driver_via_api()
 
         mock_cmd.assert_called_once()
 
@@ -1708,7 +1708,7 @@ def test_poll_k8s_driver_raises_after_consecutive_api_errors(self, mock_get_clie
         mock_client.read_namespaced_pod.side_effect = api_error
 
         with pytest.raises(RuntimeError, match="K8s API unreachable"):
-            hook._poll_k8s_driver_via_api()
+            hook.poll_k8s_driver_via_api()
 
         assert mock_client.read_namespaced_pod.call_count == 3
 
@@ -1722,7 +1722,7 @@ def test_poll_k8s_driver_exits_cleanly_on_404(self, mock_get_client):
         mock_client = mock_get_client.return_value
         mock_client.read_namespaced_pod.side_effect = kube_client.ApiException(status=404, reason="Not Found")
 
-        hook._poll_k8s_driver_via_api()
+        hook.poll_k8s_driver_via_api()
 
         mock_client.delete_namespaced_pod.assert_not_called()
 
@@ -2214,3 +2214,21 @@ def test_on_kill_tolerates_rm_failure(self, mock_put):
         hook.on_kill()
 
         mock_put.assert_called_once()
+
+    def test_conf_mutation_writes_through(self):
+        hook = SparkSubmitHook(conn_id="")
+        hook.conf["spark.foo"] = "bar"
+        assert hook._conf["spark.foo"] == "bar"
+
+    def test_kubernetes_driver_pod_writes_through(self):
+        hook = SparkSubmitHook(conn_id="")
+        hook.kubernetes_driver_pod = "spark-app-abc-driver"
+        assert hook._kubernetes_driver_pod == "spark-app-abc-driver"
+        assert hook.kubernetes_driver_pod == "spark-app-abc-driver"
+
+    def test_yarn_application_id_is_get_only(self):
+        hook = SparkSubmitHook(conn_id="")
+        hook._yarn_application_id = "application_1486558679801_1820"
+        assert hook.yarn_application_id == "application_1486558679801_1820"
+        with pytest.raises(AttributeError):
+            hook.yarn_application_id = "application_1486558679801_1821"

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

@@ -748,8 +748,8 @@ def capture(url, timeout):
     def test_yarn_first_run_persists_app_id_before_polling(self):
         operator = self._make_operator()
         operator._hook = self._make_hook(is_yarn_cluster=True)
-        operator._hook._conf = {}
-        operator._hook._yarn_application_id = "application_1234_0001"
+        operator._hook.conf = {}
+        operator._hook.yarn_application_id = "application_1234_0001"
         operator._hook.submit.return_value = None
 
         task_store = FakeTaskStateStore()
@@ -791,8 +791,8 @@ def test_yarn_retry_skips_already_succeeded_app(self):
     def test_yarn_retry_resubmits_after_failed_app(self):
         operator = self._make_operator()
         operator._hook = self._make_hook(is_yarn_cluster=True)
-        operator._hook._conf = {}
-        operator._hook._yarn_application_id = "application_1234_0002"
+        operator._hook.conf = {}
+        operator._hook.yarn_application_id = "application_1234_0002"
         operator._hook.submit.return_value = None
         task_store = FakeTaskStateStore({"spark_job_id": "application_1234_0001"})
 
@@ -808,19 +808,19 @@ def test_yarn_retry_resubmits_after_failed_app(self):
     def test_yarn_injects_wait_app_completion_false(self):
         operator = self._make_operator()
         hook = self._make_hook(is_yarn_cluster=True)
-        hook._conf = {}
-        hook._yarn_application_id = "application_1234_0001"
+        hook.conf = {}
+        hook.yarn_application_id = "application_1234_0001"
         hook.submit.return_value = None
         operator._hook = hook
 
         operator.submit_job(context={})
 
-        assert hook._conf.get("spark.yarn.submit.waitAppCompletion") == "false"
+        assert hook.conf.get("spark.yarn.submit.waitAppCompletion") == "false"
 
     def test_yarn_raises_if_wait_app_completion_true(self):
         operator = self._make_operator()
         hook = self._make_hook(is_yarn_cluster=True)
-        hook._conf = {"spark.yarn.submit.waitAppCompletion": "true"}
+        hook.conf = {"spark.yarn.submit.waitAppCompletion": "true"}
         operator._hook = hook
 
         with pytest.raises(
@@ -869,16 +869,16 @@ def simulate_failed_tracking():
             operator.poll_until_complete("driver-001", {})
 
     def test_on_kill_sends_authenticated_kill_to_yarn_rm(self):
-        """operator.on_kill() must call _kill_yarn_application so Kerberos auth is applied."""
+        """operator.on_kill() must call kill_yarn_application so Kerberos auth is applied."""
         operator = self._make_operator()
         hook = self._make_hook(is_yarn_cluster=True)
         hook._is_yarn_cluster_mode = True
-        hook._yarn_application_id = "application_1234_0001"
+        hook.yarn_application_id = "application_1234_0001"
         operator._hook = hook
 
         operator.on_kill()
 
-        hook._kill_yarn_application.assert_called_once_with("application_1234_0001")
+        hook.kill_yarn_application.assert_called_once_with("application_1234_0001")
 
     def test_yarn_cluster_reconnect_without_rm_api_raises(self):
         """durable=True + yarn_track_via_rm_api=False must raise - RM API is required for resume."""
@@ -917,7 +917,7 @@ def _make_k8s_hook(self):
         hook._is_kubernetes = True
         hook._is_yarn = False
         hook._is_yarn_cluster_mode = False
-        hook._conf = {}
+        hook.conf = {}
         return hook
 
     def test_execute_calls_submit_then_poll_when_flag_set(self):
@@ -926,12 +926,12 @@ def test_execute_calls_submit_then_poll_when_flag_set(self):
         operator._hook = hook
         call_order = []
         hook.submit.side_effect = lambda *a, **kw: call_order.append("submit")
-        hook._poll_k8s_driver_via_api.side_effect = lambda: call_order.append("poll")
+        hook.poll_k8s_driver_via_api.side_effect = lambda: call_order.append("poll")
 
         operator.execute(context={})
 
         hook.submit.assert_called_once_with("test.jar")
-        hook._poll_k8s_driver_via_api.assert_called_once()
+        hook.poll_k8s_driver_via_api.assert_called_once()
         assert call_order == ["submit", "poll"]
 
     def test_execute_falls_through_to_plain_submit_when_flag_off(self):
@@ -944,25 +944,25 @@ def test_execute_falls_through_to_plain_submit_when_flag_off(self):
         operator.execute(context={})
 
         hook.submit.assert_called_once_with("test.jar")
-        hook._poll_k8s_driver_via_api.assert_not_called()
+        hook.poll_k8s_driver_via_api.assert_not_called()
 
     def test_k8s_submit_job_returns_encoded_external_id(self):
         operator = self._make_operator(track_driver_via_k8s_api=True)
         hook = self._make_k8s_hook()
-        hook._kubernetes_driver_pod = "spark-abc-driver"
+        hook.kubernetes_driver_pod = "spark-abc-driver"
         hook._connection = {"namespace": "mynamespace"}
         operator._hook = hook
 
         result = operator.submit_job(context={})
 
         assert result == "mynamespace:spark-abc-driver"
-        assert hook._conf.get("spark.kubernetes.submission.waitAppCompletion") == "false"
+        assert hook.conf.get("spark.kubernetes.submission.waitAppCompletion") == "false"
         hook.submit.assert_called_once_with("test.jar")
 
     def test_k8s_submit_job_raises_when_pod_name_missing(self):
         operator = self._make_operator(track_driver_via_k8s_api=True)
         hook = self._make_k8s_hook()
-        hook._kubernetes_driver_pod = None
+        hook.kubernetes_driver_pod = None
         hook._connection = {"namespace": "mynamespace"}
         operator._hook = hook
 
@@ -1043,14 +1043,14 @@ def test_k8s_poll_until_complete_sets_pod_name_and_calls_poll_api(self):
 
         operator.poll_until_complete("mynamespace:spark-abc-driver", {})
 
-        assert hook._kubernetes_driver_pod == "spark-abc-driver"
-        hook._poll_k8s_driver_via_api.assert_called_once()
+        assert hook.kubernetes_driver_pod == "spark-abc-driver"
+        hook.poll_k8s_driver_via_api.assert_called_once()
 
     @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task_state_store requires Airflow 3.3+")
     def test_k8s_poll_until_complete_writes_succeeded_to_task_store(self):
         operator = self._make_operator(track_driver_via_k8s_api=True)
         hook = self._make_k8s_hook()
-        hook._poll_k8s_driver_via_api.return_value = "Succeeded"
+        hook.poll_k8s_driver_via_api.return_value = "Succeeded"
         operator._hook = hook
         task_store = FakeTaskStateStore()
 
@@ -1062,7 +1062,7 @@ def test_k8s_poll_until_complete_writes_succeeded_to_task_store(self):
     def test_k8s_polling_does_not_write_task_store_when_reconnect_disabled(self):
         operator = self._make_operator(track_driver_via_k8s_api=True, durable=False)
         hook = self._make_k8s_hook()
-        hook._poll_k8s_driver_via_api.return_value = "Succeeded"
+        hook.poll_k8s_driver_via_api.return_value = "Succeeded"
         operator._hook = hook
         task_store = FakeTaskStateStore()
 
@@ -1073,7 +1073,7 @@ def test_k8s_polling_does_not_write_task_store_when_reconnect_disabled(self):
     def test_k8s_poll_until_complete_does_not_cache_and_reraises_on_failure(self):
         operator = self._make_operator(track_driver_via_k8s_api=True)
         hook = self._make_k8s_hook()
-        hook._poll_k8s_driver_via_api.side_effect = RuntimeError("Spark application failed (phase=Failed)")
+        hook.poll_k8s_driver_via_api.side_effect = RuntimeError("Spark application failed (phase=Failed)")
         operator._hook = hook
         task_store = FakeTaskStateStore()
 
@@ -1096,7 +1096,7 @@ def test_k8s_execute_persists_pod_id_when_durable(self):
         """execute() with durable=True stores the pod ID in task_store before polling."""
         operator = self._make_operator(track_driver_via_k8s_api=True, durable=True)
         hook = self._make_k8s_hook()
-        hook._kubernetes_driver_pod = "spark-abc-driver"
+        hook.kubernetes_driver_pod = "spark-abc-driver"
         hook._connection = {"namespace": "mynamespace"}
         operator._hook = hook
         task_store = FakeTaskStateStore()
@@ -1119,7 +1119,7 @@ def test_k8s_execute_durable_false_does_not_persist_pod_id(self):
         """execute() with durable=False does not write spark_job_id to task_store."""
         operator = self._make_operator(track_driver_via_k8s_api=True, durable=False)
         hook = self._make_k8s_hook()
-        hook._kubernetes_driver_pod = "spark-abc-driver"
+        hook.kubernetes_driver_pod = "spark-abc-driver"
         hook._connection = {"namespace": "mynamespace"}
         operator._hook = hook
         task_store = FakeTaskStateStore()