Fix SparkPipelinesHook for Spark Connect connections

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

@@ -17,11 +17,13 @@
 # under the License.
 from __future__ import annotations
 
+import os
 import subprocess
+import sys
 from typing import Any
 
-from airflow.providers.apache.spark.hooks.spark_submit import SparkSubmitHook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.apache.spark.hooks.spark_submit import DEFAULT_SPARK_BINARY, SparkSubmitHook
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
 
 
 class SparkPipelinesException(AirflowException):
@@ -35,6 +37,20 @@ class SparkPipelinesHook(SparkSubmitHook):
     Extends SparkSubmitHook to leverage existing connection management while providing
     pipeline-specific functionality.
 
+    Two connection modes are supported:
+
+    * Legacy spark-submit-style (``spark`` / ``yarn`` / ``k8s`` connection types) —
+      invokes the ``spark-pipelines`` launcher with ``--master``, ``--deploy-mode``
+      and the rest of the standard cluster-manager flags assembled by
+      :class:`~airflow.providers.apache.spark.hooks.spark_submit.SparkSubmitHook`.
+    * Spark Connect (``spark_connect`` connection type, Spark 4.x+) — sets
+      ``SPARK_REMOTE`` from the connection's ``sc://`` URI and invokes the
+      Connect-native ``pyspark.pipelines.cli`` Python module directly. The
+      cluster-manager flags are *not* emitted: the Connect-native CLI rejects
+      them with ``SparkException: Remote cannot be specified with master and/or
+      deploy mode``, and the ``spark-pipelines`` bash launcher itself starts a
+      JVM ``SparkContext`` that collides with the Connect daemon's gRPC port.
+
     :param pipeline_spec: Path to the pipeline specification file (YAML)
     :param pipeline_command: The spark-pipelines command to run ('run', 'dry-run')
     """
@@ -52,7 +68,56 @@ def __init__(
         if pipeline_command not in ["run", "dry-run"]:
             raise ValueError(f"Invalid pipeline command: {pipeline_command}. Must be 'run' or 'dry-run'")
 
+    def _resolve_connection(self) -> dict[str, Any]:
+        """
+        Resolve the configured connection, branching on Spark Connect.
+
+        For ``spark_connect``-typed connections, populate a ``spark_remote`` key
+        with the ``sc://`` URI rendered by
+        :class:`~airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook`
+        and zero out the spark-submit cluster-manager fields — the Connect-native
+        CLI doesn't consume them. For every other connection type, defer to the
+        parent's resolver so spark-submit-style behaviour is unchanged.
+        """
+        try:
+            conn = self.get_connection(self._conn_id)
+        except AirflowNotFoundException:
+            # No connection configured — fall through to spark-submit defaults.
+            return super()._resolve_connection()
+
+        if conn.conn_type != "spark_connect":
+            return super()._resolve_connection()
+
+        # Local import: SparkConnectHook lives in the same provider but loading
+        # it eagerly would create an unnecessary cycle for spark-submit-only
+        # deployments that never touch Connect.
+        from airflow.providers.apache.spark.hooks.spark_connect import SparkConnectHook
+
+        return {
+            # ``master`` is consumed by the parent's ``__init__`` (substring
+            # checks for ``yarn``/``k8s``/``spark://``); leave it empty so none
+            # of those branches match. Connect mode never calls
+            # ``_build_spark_common_args``, which is the only thing that would
+            # actually emit ``--master`` to the CLI.
+            "master": "",
+            "queue": None,
+            "deploy_mode": None,
+            "spark_binary": self.spark_binary or DEFAULT_SPARK_BINARY,
+            "namespace": None,
+            "principal": self._principal,
+            "keytab": self._keytab,
+            "spark_remote": SparkConnectHook(conn_id=self._conn_id).get_connection_url(),
+        }
+
     def _get_spark_binary_path(self) -> list[str]:
+        if self._connection.get("spark_remote"):
+            # The ``spark-pipelines`` bash launcher routes through
+            # ``spark-class`` → ``SparkSubmit`` → JVM ``SparkContext``, which
+            # appends cluster-manager flags the Connect-native CLI rejects and
+            # binds an in-process Connect server that collides with the
+            # long-running daemon. Invoke the underlying Python module directly
+            # so ``SparkSession.builder.getOrCreate()`` becomes a Connect client.
+            return [sys.executable, "-m", "pyspark.pipelines.cli"]
         return ["spark-pipelines"]
 
     def _build_spark_pipelines_command(self) -> list[str]:
@@ -61,16 +126,17 @@ def _build_spark_pipelines_command(self) -> list[str]:
 
         :return: full command to be executed
         """
-        # Start with spark-pipelines binary and command
         connection_cmd = self._get_spark_binary_path()
         connection_cmd.append(self.pipeline_command)
 
-        # Add pipeline spec if provided
         if self.pipeline_spec:
             connection_cmd.extend(["--spec", self.pipeline_spec])
 
-        # Reuse parent's common spark argument building logic
-        connection_cmd.extend(self._build_spark_common_args())
+        if not self._connection.get("spark_remote"):
+            # Reuse parent's common spark argument building logic. The
+            # Connect-native CLI rejects --master/--deploy-mode/--name, so only
+            # emit them for legacy spark-submit-style connections.
+            connection_cmd.extend(self._build_spark_common_args())
 
         self.log.info("Spark-Pipelines cmd: %s", self._mask_cmd(connection_cmd))
         return connection_cmd
@@ -82,12 +148,20 @@ def submit_pipeline(self, **kwargs: Any) -> None:
         :param kwargs: extra arguments to Popen (see subprocess.Popen)
         """
         pipelines_cmd = self._build_spark_pipelines_command()
-
-        if self._env:
-            import os
-
+        spark_remote = self._connection.get("spark_remote")
+
+        # ``self._env`` is only populated by ``_build_spark_common_args`` —
+        # which the Connect path skips — so fall back to ``self._env_vars``
+        # (the operator-supplied env_vars kwarg) for that path.
+        env_overrides: dict[str, str] = dict(self._env or self._env_vars or {})
+        if spark_remote:
+            # Don't clobber a SPARK_REMOTE the operator caller already set via
+            # env_vars; that takes precedence for failover routing.
+            env_overrides.setdefault("SPARK_REMOTE", spark_remote)
+
+        if env_overrides:
             env = os.environ.copy()
-            env.update(self._env)
+            env.update(env_overrides)
             kwargs["env"] = env
 
         self._submit_sp = subprocess.Popen(

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

@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+import sys
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -183,3 +184,132 @@ def test_submit_calls_submit_pipeline(self):
         with patch.object(self.hook, "submit_pipeline") as mock_submit_pipeline:
             self.hook.submit("dummy_application")
             mock_submit_pipeline.assert_called_once()
+
+    # ------------------------------------------------------------------ #
+    # Spark Connect path: spark_connect-typed conn_id should bypass the   #
+    # spark-submit launcher and the cluster-manager flags entirely, and   #
+    # set SPARK_REMOTE in the subprocess environment instead.             #
+    # ------------------------------------------------------------------ #
+
+    @patch(
+        "airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook.get_connection_url",
+        return_value="sc://spark-connect.example:15002/",
+    )
+    @patch("airflow.providers.apache.spark.hooks.spark_pipelines.SparkPipelinesHook.get_connection")
+    def test_build_pipelines_command_spark_connect_skips_cluster_args(
+        self, mock_get_connection, mock_get_url
+    ):
+        spark_connect_conn = MagicMock()
+        spark_connect_conn.conn_type = "spark_connect"
+        spark_connect_conn.host = "spark-connect.example"
+        spark_connect_conn.port = 15002
+        spark_connect_conn.login = None
+        spark_connect_conn.password = None
+        # SparkSubmitHook.__init__ -> _resolve_connection() reads
+        # extra_dejson.get("queue"|"deploy-mode"|"spark-binary"|"namespace");
+        # an empty dict makes it default-and-skip cleanly.
+        spark_connect_conn.extra_dejson = {}
+        mock_get_connection.return_value = spark_connect_conn
+
+        hook = SparkPipelinesHook(
+            pipeline_spec="test_pipeline.yml",
+            pipeline_command="run",
+            conn_id="spark_connect_default",
+            num_executors=2,
+            executor_memory="2G",
+            deploy_mode="client",
+        )
+
+        cmd = hook._build_spark_pipelines_command()
+
+        # Connect-native CLI is invoked via the python module — bypassing the
+        # bash launcher that would otherwise start a colliding JVM Connect server.
+        assert cmd[:3] == [sys.executable, "-m", "pyspark.pipelines.cli"]
+        assert "run" in cmd
+        assert "--spec" in cmd
+        assert "test_pipeline.yml" in cmd
+        # Cluster-manager args MUST NOT appear: the Connect-native CLI rejects
+        # them with `Remote cannot be specified with master and/or deploy mode`.
+        assert "--master" not in cmd
+        assert "--deploy-mode" not in cmd
+        assert "--name" not in cmd
+        assert "--num-executors" not in cmd
+        assert "--executor-memory" not in cmd
+
+    @patch("subprocess.Popen")
+    @patch(
+        "airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook.get_connection_url",
+        return_value="sc://spark-connect.example:15002/",
+    )
+    @patch("airflow.providers.apache.spark.hooks.spark_pipelines.SparkPipelinesHook.get_connection")
+    def test_submit_pipeline_spark_connect_sets_spark_remote(
+        self, mock_get_connection, mock_get_url, mock_popen
+    ):
+        spark_connect_conn = MagicMock()
+        spark_connect_conn.conn_type = "spark_connect"
+        spark_connect_conn.host = "spark-connect.example"
+        spark_connect_conn.port = 15002
+        spark_connect_conn.login = None
+        spark_connect_conn.password = None
+        # SparkSubmitHook.__init__ -> _resolve_connection() reads
+        # extra_dejson.get("queue"|"deploy-mode"|"spark-binary"|"namespace");
+        # an empty dict makes it default-and-skip cleanly.
+        spark_connect_conn.extra_dejson = {}
+        mock_get_connection.return_value = spark_connect_conn
+
+        mock_process = MagicMock()
+        mock_process.wait.return_value = 0
+        mock_process.stdout = ["Run is COMPLETED."]
+        mock_popen.return_value = mock_process
+
+        hook = SparkPipelinesHook(
+            pipeline_spec="test_pipeline.yml",
+            pipeline_command="run",
+            conn_id="spark_connect_default",
+        )
+        hook.submit_pipeline()
+
+        mock_popen.assert_called_once()
+        _, popen_kwargs = mock_popen.call_args
+        env = popen_kwargs["env"]
+        assert env["SPARK_REMOTE"] == "sc://spark-connect.example:15002/"
+
+    @patch("subprocess.Popen")
+    @patch(
+        "airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook.get_connection_url",
+        return_value="sc://spark-connect.example:15002/",
+    )
+    @patch("airflow.providers.apache.spark.hooks.spark_pipelines.SparkPipelinesHook.get_connection")
+    def test_submit_pipeline_spark_connect_preserves_caller_spark_remote(
+        self, mock_get_connection, mock_get_url, mock_popen
+    ):
+        # If the operator caller has already set SPARK_REMOTE via env_vars
+        # (e.g. routing to a different daemon for failover), the hook must
+        # not clobber it with the connection's URI.
+        spark_connect_conn = MagicMock()
+        spark_connect_conn.conn_type = "spark_connect"
+        spark_connect_conn.host = "spark-connect.example"
+        spark_connect_conn.port = 15002
+        spark_connect_conn.login = None
+        spark_connect_conn.password = None
+        # SparkSubmitHook.__init__ -> _resolve_connection() reads
+        # extra_dejson.get("queue"|"deploy-mode"|"spark-binary"|"namespace");
+        # an empty dict makes it default-and-skip cleanly.
+        spark_connect_conn.extra_dejson = {}
+        mock_get_connection.return_value = spark_connect_conn
+
+        mock_process = MagicMock()
+        mock_process.wait.return_value = 0
+        mock_process.stdout = []
+        mock_popen.return_value = mock_process
+
+        hook = SparkPipelinesHook(
+            pipeline_spec="test_pipeline.yml",
+            pipeline_command="run",
+            conn_id="spark_connect_default",
+            env_vars={"SPARK_REMOTE": "sc://override.example:15002/"},
+        )
+        hook.submit_pipeline()
+
+        _, popen_kwargs = mock_popen.call_args
+        assert popen_kwargs["env"]["SPARK_REMOTE"] == "sc://override.example:15002/"