Add support for Spark Connect to pyspark decorator (#35665)

* Add support for Spark Connect to pyspark decorator

In Apache Spark 3.4 Spark Connect was introduced which
allows remote connectivity to remote Spark Cluster using
the DataFrame API.

Author: bolkedebruin

airflow/providers/apache/spark/decorators/pyspark.py

@@ -23,6 +23,7 @@
 from airflow.decorators.base import DecoratedOperator, TaskDecorator, task_decorator_factory
 from airflow.hooks.base import BaseHook
 from airflow.operators.python import PythonOperator
+from airflow.providers.apache.spark.hooks.spark_connect import SparkConnectHook
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -73,34 +74,44 @@ def execute(self, context: Context):
         from pyspark import SparkConf
         from pyspark.sql import SparkSession
 
-        conf = SparkConf().setAppName(f"{self.dag_id}-{self.task_id}")
+        conf = SparkConf()
+        conf.set("spark.app.name", f"{self.dag_id}-{self.task_id}")
 
-        master = "local[*]"
+        url = "local[*]"
         if self.conn_id:
+            # we handle both spark connect and spark standalone
             conn = BaseHook.get_connection(self.conn_id)
-            if conn.port:
-                master = f"{conn.host}:{conn.port}"
+            if conn.conn_type == SparkConnectHook.conn_type:
+                url = SparkConnectHook(self.conn_id).get_connection_url()
+            elif conn.port:
+                url = f"{conn.host}:{conn.port}"
             elif conn.host:
-                master = conn.host
+                url = conn.host
 
             for key, value in conn.extra_dejson.items():
                 conf.set(key, value)
 
-        conf.setMaster(master)
+        # you cannot have both remote and master
+        if url.startswith("sc://"):
+            conf.set("spark.remote", url)
 
         # task can override connection config
         for key, value in self.config_kwargs.items():
             conf.set(key, value)
 
+        if not conf.get("spark.remote") and not conf.get("spark.master"):
+            conf.set("spark.master", url)
+
         spark = SparkSession.builder.config(conf=conf).getOrCreate()
-        sc = spark.sparkContext
 
         if not self.op_kwargs:
             self.op_kwargs = {}
 
         op_kwargs: dict[str, Any] = dict(self.op_kwargs)
         op_kwargs["spark"] = spark
-        op_kwargs["sc"] = sc
+
+        # spark context is not available when using spark connect
+        op_kwargs["sc"] = spark.sparkContext if not conf.get("spark.remote") else None
 
         self.op_kwargs = op_kwargs
         return super().execute(context)

airflow/providers/apache/spark/hooks/spark_connect.py

@@ -0,0 +1,99 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import quote, urlparse, urlunparse
+
+from airflow.hooks.base import BaseHook
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+
+class SparkConnectHook(BaseHook, LoggingMixin):
+    """Hook for Spark Connect."""
+
+    # from pyspark's ChannelBuilder
+    PARAM_USE_SSL = "use_ssl"
+    PARAM_TOKEN = "token"
+    PARAM_USER_ID = "user_id"
+
+    conn_name_attr = "conn_id"
+    default_conn_name = "spark_connect_default"
+    conn_type = "spark_connect"
+    hook_name = "Spark Connect"
+
+    @staticmethod
+    def get_ui_field_behaviour() -> dict[str, Any]:
+        """Return custom field behaviour."""
+        return {
+            "hidden_fields": [
+                "schema",
+            ],
+            "relabeling": {"password": "Token", "login": "User ID"},
+        }
+
+    @staticmethod
+    def get_connection_form_widgets() -> dict[str, Any]:
+        """Returns connection widgets to add to connection form."""
+        from flask_babel import lazy_gettext
+        from wtforms import BooleanField
+
+        return {
+            SparkConnectHook.PARAM_USE_SSL: BooleanField(lazy_gettext("Use SSL"), default=False),
+        }
+
+    def __init__(self, conn_id: str = default_conn_name) -> None:
+        super().__init__()
+        self._conn_id = conn_id
+
+    def get_connection_url(self) -> str:
+        conn = self.get_connection(self._conn_id)
+
+        host = conn.host
+        if conn.host.find("://") == -1:
+            host = f"sc://{conn.host}"
+        if conn.port:
+            host = f"{conn.host}:{conn.port}"
+
+        url = urlparse(host)
+
+        if url.path:
+            raise ValueError("Path {url.path} is not supported in Spark Connect connection URL")
+
+        params = []
+
+        if conn.login:
+            params.append(f"{SparkConnectHook.PARAM_USER_ID}={quote(conn.login)}")
+
+        if conn.password:
+            params.append(f"{SparkConnectHook.PARAM_TOKEN}={quote(conn.password)}")
+
+        use_ssl = conn.extra_dejson.get(SparkConnectHook.PARAM_USE_SSL)
+        if use_ssl is not None:
+            params.append(f"{SparkConnectHook.PARAM_USE_SSL}={quote(str(use_ssl))}")
+
+        return urlunparse(
+            (
+                "sc",
+                url.netloc,
+                "/",
+                ";".join(params),  # params
+                "",
+                url.fragment,
+            )
+        )

airflow/providers/apache/spark/provider.yaml

@@ -51,6 +51,7 @@ versions:
 dependencies:
   - apache-airflow>=2.5.0
   - pyspark
+  - grpcio-status
 
 integrations:
   - integration-name: Apache Spark
@@ -70,13 +71,16 @@ operators:
 hooks:
   - integration-name: Apache Spark
     python-modules:
+      - airflow.providers.apache.spark.hooks.spark_connect
       - airflow.providers.apache.spark.hooks.spark_jdbc
       - airflow.providers.apache.spark.hooks.spark_jdbc_script
       - airflow.providers.apache.spark.hooks.spark_sql
       - airflow.providers.apache.spark.hooks.spark_submit
 
 
 connection-types:
+  - hook-class-name: airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook
+    connection-type: spark_connect
   - hook-class-name: airflow.providers.apache.spark.hooks.spark_jdbc.SparkJDBCHook
     connection-type: spark_jdbc
   - hook-class-name: airflow.providers.apache.spark.hooks.spark_sql.SparkSqlHook

docs/apache-airflow-providers-apache-spark/connections/spark.rst

@@ -27,7 +27,7 @@ The Apache Spark connection type enables connection to Apache Spark.
 Default Connection IDs
 ----------------------
 
-Spark Submit and Spark JDBC hooks and operators use ``spark_default`` by default. Spark SQL hooks and operators point to ``spark_sql_default`` by default.
+Spark Submit and Spark JDBC hooks and operators use ``spark_default`` by default. Spark SQL hooks and operators point to ``spark_sql_default`` by default. The Spark Connect hook uses ``spark_connect_default`` by default.
 
 Configuring the Connection
 --------------------------
@@ -45,6 +45,15 @@ Extra (optional)
     * ``spark-binary`` - The command to use for Spark submit. Some distros may use ``spark2-submit``. Default ``spark-submit``. Only ``spark-submit``, ``spark2-submit`` or ``spark3-submit`` are allowed as value.
     * ``namespace`` - Kubernetes namespace (``spark.kubernetes.namespace``) to divide cluster resources between multiple users (via resource quota).
 
+User ID (optional, only applies to Spark Connect)
+    The user ID to authenticate with the proxy.
+
+Token (optional, only applies to Spark Connect)
+    The token to authenticate with the proxy.
+
+Use SSL (optional, only applies to Spark Connect)
+    Whether to use SSL when connecting.
+
 When specifying the connection in environment variable you should specify
 it using URI syntax.
 

docs/apache-airflow-providers-apache-spark/decorators/pyspark.rst

@@ -23,7 +23,7 @@ PySpark Decorator
 =================
 
 Python callable wrapped within the ``@task.pyspark`` decorator
-is injected with a SparkContext object.
+is injected with a SparkSession and SparkContext object if available.
 
 Parameters
 ----------
@@ -49,3 +49,24 @@ that the ``spark`` and ``sc`` objects are injected into the function.
     :dedent: 4
     :start-after: [START task_pyspark]
     :end-before: [END task_pyspark]
+
+
+Spark Connect
+-------------
+
+In `Apache Spark 3.4 <https://spark.apache.org/docs/latest/spark-connect-overview.html>`_,
+Spark Connect introduced a decoupled client-server architecture
+that allows remote connectivity to Spark clusters using the DataFrame API. Using
+Spark Connect is the preferred way in Airflow to make use of the PySpark decorator,
+because it does not require to run the Spark driver on the same host as Airflow.
+To make use of Spark Connect, you prepend your host url with ``sc://``. For example,
+``sc://spark-cluster:15002``.
+
+
+Authentication
+^^^^^^^^^^^^^^
+
+Spark Connect does not have built-in authentication. The gRPC HTTP/2 interface however
+allows the use of authentication to communicate with the Spark Connect server through
+authenticating proxies. To make use of authentication make sure to create a ``Spark Connect``
+connection and set the right credentials.

generated/provider_dependencies.json

@@ -190,6 +190,7 @@
   "apache.spark": {
     "deps": [
       "apache-airflow>=2.5.0",
+      "grpcio-status",
       "pyspark"
     ],
     "cross-providers-deps": [

tests/providers/apache/spark/decorators/test_pyspark.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+from typing import Any
 from unittest import mock
 
 import pytest
@@ -27,6 +28,22 @@
 DEFAULT_DATE = timezone.datetime(2021, 9, 1)
 
 
+class FakeConfig:
+    data: dict[str, Any]
+
+    def __init__(self, data: dict[str, Any] | None = None):
+        if data:
+            self.data = data
+        else:
+            self.data = {}
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self.data.get(key, default)
+
+    def set(self, key: str, value: Any) -> None:
+        self.data[key] = value
+
+
 class TestPysparkDecorator:
     def setup_method(self):
         db.merge_conn(
@@ -38,14 +55,47 @@ def setup_method(self):
             )
         )
 
+        db.merge_conn(
+            Connection(
+                conn_id="spark-connect",
+                conn_type="spark",
+                host="sc://localhost",
+                extra="",
+            )
+        )
+
+        db.merge_conn(
+            Connection(
+                conn_id="spark-connect-auth",
+                conn_type="spark_connect",
+                host="sc://localhost",
+                password="1234",
+                login="connect",
+                extra={
+                    "use_ssl": True,
+                },
+            )
+        )
+
     @pytest.mark.db_test
-    @mock.patch("pyspark.SparkConf.setAppName")
+    @mock.patch("pyspark.SparkConf")
     @mock.patch("pyspark.sql.SparkSession")
     def test_pyspark_decorator_with_connection(self, spark_mock, conf_mock, dag_maker):
+        config = FakeConfig()
+
+        builder = mock.MagicMock()
+        spark_mock.builder.config.return_value = builder
+        builder.getOrCreate.return_value = builder
+        builder.sparkContext.return_value = builder
+
+        conf_mock.return_value = config
+
         @task.pyspark(conn_id="pyspark_local", config_kwargs={"spark.executor.memory": "2g"})
         def f(spark, sc):
             import random
 
+            assert spark is not None
+            assert sc is not None
             return [random.random() for _ in range(100)]
 
         with dag_maker():
@@ -55,14 +105,20 @@ def f(spark, sc):
         ret.operator.run(start_date=dr.execution_date, end_date=dr.execution_date)
         ti = dr.get_task_instances()[0]
         assert len(ti.xcom_pull()) == 100
-        conf_mock().set.assert_called_with("spark.executor.memory", "2g")
-        conf_mock().setMaster.assert_called_once_with("spark://none")
+        assert config.get("spark.master") == "spark://none"
+        assert config.get("spark.executor.memory") == "2g"
+        assert config.get("spark.remote") is None
+        assert config.get("spark.app.name")
+
         spark_mock.builder.config.assert_called_once_with(conf=conf_mock())
 
     @pytest.mark.db_test
-    @mock.patch("pyspark.SparkConf.setAppName")
+    @mock.patch("pyspark.SparkConf")
     @mock.patch("pyspark.sql.SparkSession")
     def test_simple_pyspark_decorator(self, spark_mock, conf_mock, dag_maker):
+        config = FakeConfig()
+        conf_mock.return_value = config
+
         e = 2
 
         @task.pyspark
@@ -76,5 +132,56 @@ def f():
         ret.operator.run(start_date=dr.execution_date, end_date=dr.execution_date)
         ti = dr.get_task_instances()[0]
         assert ti.xcom_pull() == e
-        conf_mock().setMaster.assert_called_once_with("local[*]")
+        assert config.get("spark.master") == "local[*]"
         spark_mock.builder.config.assert_called_once_with(conf=conf_mock())
+
+    @pytest.mark.db_test
+    @mock.patch("pyspark.SparkConf")
+    @mock.patch("pyspark.sql.SparkSession")
+    def test_spark_connect(self, spark_mock, conf_mock, dag_maker):
+        config = FakeConfig()
+        conf_mock.return_value = config
+
+        @task.pyspark(conn_id="spark-connect")
+        def f(spark, sc):
+            assert spark is not None
+            assert sc is None
+
+            return True
+
+        with dag_maker():
+            ret = f()
+
+        dr = dag_maker.create_dagrun()
+        ret.operator.run(start_date=dr.execution_date, end_date=dr.execution_date)
+        ti = dr.get_task_instances()[0]
+        assert ti.xcom_pull()
+        assert config.get("spark.remote") == "sc://localhost"
+        assert config.get("spark.master") is None
+        assert config.get("spark.app.name")
+        spark_mock.builder.config.assert_called_once_with(conf=conf_mock())
+
+    @pytest.mark.db_test
+    @mock.patch("pyspark.SparkConf")
+    @mock.patch("pyspark.sql.SparkSession")
+    def test_spark_connect_auth(self, spark_mock, conf_mock, dag_maker):
+        config = FakeConfig()
+        conf_mock.return_value = config
+
+        @task.pyspark(conn_id="spark-connect-auth")
+        def f(spark, sc):
+            assert spark is not None
+            assert sc is None
+
+            return True
+
+        with dag_maker():
+            ret = f()
+
+        dr = dag_maker.create_dagrun()
+        ret.operator.run(start_date=dr.execution_date, end_date=dr.execution_date)
+        ti = dr.get_task_instances()[0]
+        assert ti.xcom_pull()
+        assert config.get("spark.remote") == "sc://localhost/;user_id=connect;token=1234;use_ssl=True"
+        assert config.get("spark.master") is None
+        assert config.get("spark.app.name")

tests/providers/apache/spark/hooks/test_spark_connect.py

@@ -0,0 +1,72 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import pytest
+
+from airflow.models import Connection
+from airflow.providers.apache.spark.hooks.spark_connect import SparkConnectHook
+from airflow.utils import db
+
+pytestmark = pytest.mark.db_test
+
+
+class TestSparkConnectHook:
+    def setup_method(self):
+        db.merge_conn(
+            Connection(
+                conn_id="spark-default",
+                conn_type="spark_connect",
+                host="sc://spark-host",
+                port=1000,
+                login="spark-user",
+                password="1234",
+                extra='{"queue": "root.etl", "deploy-mode": "cluster"}',
+            )
+        )
+
+        db.merge_conn(
+            Connection(
+                conn_id="spark-test",
+                conn_type="spark_connect",
+                host="nowhere",
+                login="spark-user",
+            )
+        )
+
+        db.merge_conn(
+            Connection(
+                conn_id="spark-app",
+                conn_type="spark_connect",
+                host="sc://cluster/app",
+                login="spark-user",
+            )
+        )
+
+    def test_get_connection_url(self):
+        expected_url = "sc://spark-host:1000/;user_id=spark-user;token=1234"
+        hook = SparkConnectHook(conn_id="spark-default")
+        assert hook.get_connection_url() == expected_url
+
+        expected_url = "sc://nowhere/;user_id=spark-user"
+        hook = SparkConnectHook(conn_id="spark-test")
+        assert hook.get_connection_url() == expected_url
+
+        hook = SparkConnectHook(conn_id="spark-app")
+        with pytest.raises(ValueError):
+            hook.get_connection_url()