fix spark-kubernetes-operator compatibility

airflow/providers/cncf/kubernetes/example_dags/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-sensor.py

@@ -0,0 +1,31 @@
+#
+# 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.
+"""Example of a DAG that uses :py:class:`SparkKubernetesOperator` with `watch` that is responsible
+for the spark job states and logs.
+"""
+from __future__ import annotations
+
+import datetime
+
+from airflow import DAG
+from airflow.providers.cncf.kubernetes.operators.spark_kubernetes import SparkKubernetesOperator
+
+with DAG("example_spark_k8s_operator_with_watch", start_date=datetime.datetime(2023, 1, 1)) as dag:
+    SparkKubernetesOperator(
+        task_id="spark-operator", application_file="spark-pi.yaml", namespace="spark-apps", watch=True
+    )

airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-watch.py

@@ -0,0 +1,43 @@
+#
+# 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.
+"""Example of a DAG that uses :py:class:`SparkKubernetesOperator` with sensor.
+The operator creates the spark job on the kubernetes and the sensor watches the state of the job.
+"""
+from __future__ import annotations
+
+import datetime
+
+from airflow import DAG
+from airflow.providers.cncf.kubernetes.operators.spark_kubernetes import SparkKubernetesOperator
+from airflow.providers.cncf.kubernetes.sensors.spark_kubernetes import SparkKubernetesSensor
+
+with DAG("example_spark_k8s_with_sensor", start_date=datetime.datetime(2023, 1, 1)) as dag:
+    t1 = SparkKubernetesOperator(
+        task_id="spark_pi_submit",
+        namespace="default",
+        application_file="spark-pi.yaml",
+        do_xcom_push=True,
+    )
+
+    t2 = SparkKubernetesSensor(
+        task_id="spark_pi_sensor",
+        namespace="default",
+        application_name="{{ task_instance.xcom_pull(task_ids='spark_pi_submit')['metadata']['name'] }}",
+    )
+
+    t1 >> t2

airflow/providers/cncf/kubernetes/example_dags/spark-pi.yaml

@@ -0,0 +1,55 @@
+# 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.
+---
+apiVersion: "sparkoperator.k8s.io/v1beta2"
+kind: SparkApplication
+metadata:
+  name: spark-pi
+spec:
+  type: Scala
+  mode: cluster
+  image: "gcr.io/spark-operator/spark:v3.1.1"
+  imagePullPolicy: Always
+  mainClass: org.apache.spark.examples.SparkPi
+  mainApplicationFile: "local:///opt/spark/examples/jars/spark-examples_2.12-3.1.1.jar"
+  sparkVersion: "3.1.1"
+  restartPolicy:
+    type: Never
+  volumes:
+    - name: "test-volume"
+      hostPath:
+        path: "/tmp"
+        type: Directory
+  driver:
+    cores: 1
+    coreLimit: "1200m"
+    memory: "512m"
+    labels:
+      version: 3.1.1
+    serviceAccount: spark
+    volumeMounts:
+      - name: "test-volume"
+        mountPath: "/tmp"
+  executor:
+    cores: 1
+    instances: 1
+    memory: "512m"
+    labels:
+      version: 3.1.1
+    volumeMounts:
+      - name: "test-volume"
+        mountPath: "/tmp"

airflow/providers/cncf/kubernetes/hooks/kubernetes.py

@@ -277,7 +277,13 @@ def custom_object_client(self) -> client.CustomObjectsApi:
         return client.CustomObjectsApi(api_client=self.api_client)
 
     def create_custom_object(
-        self, group: str, version: str, plural: str, body: str | dict, namespace: str | None = None
+        self,
+        group: str,
+        version: str,
+        plural: str,
+        body: str | dict,
+        namespace: str | None = None,
+        delete_existing_object: bool = True,
     ):
         """
         Creates custom resource definition object in Kubernetes.
@@ -287,6 +293,7 @@ def create_custom_object(
         :param plural: api plural
         :param body: crd object definition
         :param namespace: kubernetes namespace
+        :param delete_existing_object: whether to delete the existing object before creating or not
         """
         api: client.CustomObjectsApi = self.custom_object_client
 
@@ -295,6 +302,20 @@ def create_custom_object(
         else:
             body_dict = body
 
+        if delete_existing_object and "name" in body_dict["metadata"]:
+            try:
+                api.delete_namespaced_custom_object(
+                    group=group,
+                    version=version,
+                    namespace=namespace,
+                    plural=plural,
+                    name=body_dict["metadata"]["name"],
+                )
+
+                self.log.warning("Deleted %s with the same name", plural)
+            except client.ApiException:
+                self.log.info("%s %s not found", plural, body_dict["metadata"]["name"])
+
         response = api.create_namespaced_custom_object(
             group=group,
             version=version,

airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py

@@ -46,6 +46,7 @@ class SparkKubernetesOperator(BaseOperator):
     :param api_group: kubernetes api group of sparkApplication
     :param api_version: kubernetes api version of sparkApplication
     :param watch: whether to watch the job status and logs or not
+    :param delete_existing_object: whether to delete the existing object before creating or not
     """
 
     template_fields: Sequence[str] = ("application_file", "namespace")
@@ -64,6 +65,7 @@ def __init__(
         cluster_context: str | None = None,
         config_file: str | None = None,
         watch: bool = False,
+        delete_existing_object: bool = True,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -77,6 +79,7 @@ def __init__(
         self.cluster_context = cluster_context
         self.config_file = config_file
         self.watch = watch
+        self.delete_existing_object = delete_existing_object
 
         self.hook = KubernetesHook(
             conn_id=self.kubernetes_conn_id,
@@ -109,6 +112,7 @@ def execute(self, context: Context):
                     namespace=namespace,
                 )
                 is_job_created = True
+                is_deleted = self.delete_existing_object
                 for event in namespace_event_stream:
                     obj = event["object"]
                     if event["object"].last_timestamp >= datetime.datetime.strptime(
@@ -129,6 +133,9 @@ def execute(self, context: Context):
                             "SparkApplicationFailed",
                             "SparkApplicationDeleted",
                         ]:
+                            if is_deleted and obj.reason == "SparkApplicationDeleted":
+                                is_deleted = False
+                                continue
                             is_job_created = False
                             raise AirflowException(obj.message)
                         elif obj.reason == "SparkApplicationCompleted":

tests/providers/cncf/kubernetes/hooks/test_kubernetes.py

@@ -373,6 +373,37 @@ def test_missing_default_connection_is_ok(self, remove_default_conn):
         with pytest.raises(AirflowNotFoundException, match="The conn_id `some_conn` isn't defined"):
             hook.conn_extras
 
+    @patch("kubernetes.config.kube_config.KubeConfigLoader")
+    @patch("kubernetes.config.kube_config.KubeConfigMerger")
+    @patch(f"{HOOK_MODULE}.client.CustomObjectsApi")
+    def test_create_custom_object(
+        self, mock_custom_object_api, mock_kube_config_merger, mock_kube_config_loader
+    ):
+        hook = KubernetesHook()
+        hook.create_custom_object(
+            group="group",
+            version="version",
+            plural="plural",
+            namespace="namespace",
+            body={"metadata": {"name": "name"}},
+        )
+
+        mock_custom_object_api.return_value.delete_namespaced_custom_object.assert_called_once_with(
+            group="group",
+            version="version",
+            plural="plural",
+            name="name",
+            namespace="namespace",
+        )
+
+        mock_custom_object_api.return_value.create_namespaced_custom_object.assert_called_once_with(
+            group="group",
+            version="version",
+            plural="plural",
+            namespace="namespace",
+            body={"metadata": {"name": "name"}},
+        )
+
     @patch("kubernetes.config.kube_config.KubeConfigLoader")
     @patch("kubernetes.config.kube_config.KubeConfigMerger")
     @patch(f"{HOOK_MODULE}.client.CustomObjectsApi")