From d1bcee3bab2cad54855957f97aa44164cadf653c Mon Sep 17 00:00:00 2001
From: Hanna Liashchuk <g.liashchuk@temabit.com>
Date: Wed, 21 Sep 2022 18:08:51 +0300
Subject: [PATCH 1/4]  [ISSUE-18468] add container_name  option for
 SparkKubernetesSensor

---
 .../providers/cncf/kubernetes/sensors/spark_kubernetes.py    | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py b/airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py
index 8cfdbf753c27b..5f5a3eb60de27 100644
--- a/airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py
+++ b/airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py
@@ -39,6 +39,7 @@ class SparkKubernetesSensor(BaseSensorOperator):
 
     :param application_name: spark Application resource name
     :param namespace: the kubernetes namespace where the sparkApplication reside in
+    :param container_name: the kubernetes container name where the sparkApplication reside in
     :param kubernetes_conn_id: The :ref:`kubernetes connection<howto/connection:kubernetes>`
         to Kubernetes cluster.
     :param attach_log: determines whether logs for driver pod should be appended to the sensor log
@@ -56,6 +57,7 @@ def __init__(
         application_name: str,
         attach_log: bool = False,
         namespace: str | None = None,
+        container_name: str = "spark-kubernetes-driver",
         kubernetes_conn_id: str = "kubernetes_default",
         api_group: str = 'sparkoperator.k8s.io',
         api_version: str = 'v1beta2',
@@ -65,6 +67,7 @@ def __init__(
         self.application_name = application_name
         self.attach_log = attach_log
         self.namespace = namespace
+        self.container_name = container_name
         self.kubernetes_conn_id = kubernetes_conn_id
         self.hook = KubernetesHook(conn_id=self.kubernetes_conn_id)
         self.api_group = api_group
@@ -84,7 +87,7 @@ def _log_driver(self, application_state: str, response: dict) -> None:
         log_method = self.log.error if application_state in self.FAILURE_STATES else self.log.info
         try:
             log = ""
-            for line in self.hook.get_pod_logs(driver_pod_name, namespace=namespace):
+            for line in self.hook.get_pod_logs(driver_pod_name, namespace=namespace, container=self.container_name):
                 log += line.decode()
             log_method(log)
         except client.rest.ApiException as e:

From bb4ee39ce21f70e132be723108a471563ce66e53 Mon Sep 17 00:00:00 2001
From: Hanna Liashchuk <g.liashchuk@temabit.com>
Date: Wed, 21 Sep 2022 18:20:18 +0300
Subject: [PATCH 2/4]  [ISSUE-18468] add test

---
 .../sensors/test_spark_kubernetes.py          | 81 +++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py b/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
index c8f6764cac3fb..4bf899331dc91 100644
--- a/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
+++ b/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
@@ -483,6 +483,63 @@
     },
 }
 
+TEST_SIDECAR_DRIVER_APPLICATION = {
+    "apiVersion": "sparkoperator.k8s.io/v1beta2",
+    "kind": "SparkApplication",
+    "metadata": {
+        "creationTimestamp": "2020-02-27T08:03:02Z",
+        "generation": 4,
+        "name": "spark-pi",
+        "namespace": "default",
+        "resourceVersion": "552073",
+        "selfLink": "/apis/sparkoperator.k8s.io/v1beta2/namespaces/default/sparkapplications/spark-pi",
+        "uid": "0c93527d-4dd9-4006-b40a-1672872e8d6f",
+    },
+    "spec": {
+        "arguments": ["100000"],
+        "driver": {
+            "coreLimit": "1200m",
+            "cores": 1,
+            "labels": {"version": "2.4.4"},
+            "memory": "512m",
+            "serviceAccount": "default",
+            "sidecars": {
+                "name": "sidecar-container",
+                "image": "busybox",
+                "command": "[\"sh\",\"-c\",\"for i in {2..10}; do echo \"output: $i\"; done\"]"
+            }
+        },
+        "executor": {
+            "cores": 1,
+            "instances": 1,
+            "labels": {"version": "2.4.4"},
+            "memory": "512m",
+        },
+        "image": "gcr.io/spark-operator/spark:v2.4.4-gcs-prometheus",
+        "imagePullPolicy": "Always",
+        "mainApplicationFile": "local:///opt/spark/examples/jars/spark-examples_2.11-2.4.4.jar",
+        "mainClass": "org.apache.spark.examples.SparkPi",
+        "mode": "cluster",
+        "monitoring": {
+            "exposeDriverMetrics": True,
+            "exposeExecutorMetrics": True,
+            "prometheus": {
+                "jmxExporterJar": "/prometheus/jmx_prometheus_javaagent-0.11.0.jar",
+                "port": 8090,
+            },
+        },
+        "restartPolicy": {"type": "Never"},
+        "sparkVersion": "2.4.4",
+        "type": "Scala",
+    },
+    "status": {
+        "applicationState": {"state": "PENDING_RERUN"},
+        "driverInfo": {},
+        "lastSubmissionAttemptTime": None,
+        "terminationTime": None,
+    },
+}
+
 TEST_POD_LOGS = [b"LOG LINE 1\n", b"LOG LINE 2"]
 TEST_POD_LOG_RESULT = "LOG LINE 1\nLOG LINE 2"
 
@@ -773,3 +830,27 @@ def test_driver_logging_error(
         )
         sensor.poke(None)
         warn_log_call.assert_called_once()
+
+    @patch(
+        "kubernetes.client.api.custom_objects_api.CustomObjectsApi.get_namespaced_custom_object",
+        return_value=TEST_SIDECAR_DRIVER_APPLICATION,
+    )
+    @patch("logging.Logger.info")
+    @patch(
+        "airflow.providers.cncf.kubernetes.hooks.kubernetes.KubernetesHook.get_pod_logs",
+        return_value=TEST_POD_LOGS,
+    )
+    def test_sidecar_driver_logging_completed(
+        self, mock_log_call, info_log_call
+    ):
+        sensor = SparkKubernetesSensor(
+            application_name="spark_pi",
+            attach_log=True,
+            dag=self.dag,
+            task_id="test_task_id",
+        )
+        sensor.poke(None)
+        mock_log_call.assert_called_once_with("spark-pi-2020-02-24-1-driver", namespace="default")
+        log_info_call = info_log_call.mock_calls[2]
+        log_value = log_info_call[1][0]
+        assert log_value == TEST_POD_LOG_RESULT

From b8f6cbdb938a3a4c3b4c778abd191d3880acdf15 Mon Sep 17 00:00:00 2001
From: Hanna Liashchuk <g.liashchuk@temabit.com>
Date: Thu, 22 Sep 2022 15:24:06 +0300
Subject: [PATCH 3/4] fix checks

---
 .../cncf/kubernetes/sensors/spark_kubernetes.py      |  4 +++-
 .../cncf/kubernetes/sensors/test_spark_kubernetes.py | 12 +++++++-----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py b/airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py
index 5f5a3eb60de27..88799202654a5 100644
--- a/airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py
+++ b/airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py
@@ -87,7 +87,9 @@ def _log_driver(self, application_state: str, response: dict) -> None:
         log_method = self.log.error if application_state in self.FAILURE_STATES else self.log.info
         try:
             log = ""
-            for line in self.hook.get_pod_logs(driver_pod_name, namespace=namespace, container=self.container_name):
+            for line in self.hook.get_pod_logs(
+                driver_pod_name, namespace=namespace, container=self.container_name
+            ):
                 log += line.decode()
             log_method(log)
         except client.rest.ApiException as e:
diff --git a/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py b/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
index 4bf899331dc91..6af3a4d4628b3 100644
--- a/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
+++ b/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
@@ -506,8 +506,8 @@
             "sidecars": {
                 "name": "sidecar-container",
                 "image": "busybox",
-                "command": "[\"sh\",\"-c\",\"for i in {2..10}; do echo \"output: $i\"; done\"]"
-            }
+                "command": "[\"sh\",\"-c\",\"for i in {2..10}; do echo \"output: $i\"; done\"]",
+            },
         },
         "executor": {
             "cores": 1,
@@ -783,7 +783,8 @@ def test_driver_logging_failure(
         )
         with pytest.raises(AirflowException):
             sensor.poke(None)
-        mock_log_call.assert_called_once_with("spark-pi-driver", namespace="default")
+        mock_log_call.assert_called_once_with("spark-pi-driver", namespace="default",
+                                              container='spark-kubernetes-driver')
         error_log_call.assert_called_once_with(TEST_POD_LOG_RESULT)
 
     @patch(
@@ -805,7 +806,8 @@ def test_driver_logging_completed(
             task_id="test_task_id",
         )
         sensor.poke(None)
-        mock_log_call.assert_called_once_with("spark-pi-2020-02-24-1-driver", namespace="default")
+        mock_log_call.assert_called_once_with("spark-pi-2020-02-24-1-driver", namespace="default",
+                                              container='spark-kubernetes-driver')
         log_info_call = info_log_call.mock_calls[2]
         log_value = log_info_call[1][0]
         assert log_value == TEST_POD_LOG_RESULT
@@ -841,7 +843,7 @@ def test_driver_logging_error(
         return_value=TEST_POD_LOGS,
     )
     def test_sidecar_driver_logging_completed(
-        self, mock_log_call, info_log_call
+        self, mock_log_call, info_log_call, mock_get_namespaced_crd, mock_kube_conn
     ):
         sensor = SparkKubernetesSensor(
             application_name="spark_pi",

From 60f35c3db45cc9620fdc60bfdd16cd76bd2e5bfe Mon Sep 17 00:00:00 2001
From: Hanna Liashchuk <g.liashchuk@temabit.com>
Date: Fri, 23 Sep 2022 17:15:51 +0300
Subject: [PATCH 4/4] fix tests

---
 .../sensors/test_spark_kubernetes.py          | 77 +++++++++++--------
 1 file changed, 43 insertions(+), 34 deletions(-)

diff --git a/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py b/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
index 6af3a4d4628b3..4f69e73a58d87 100644
--- a/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
+++ b/tests/providers/cncf/kubernetes/sensors/test_spark_kubernetes.py
@@ -483,60 +483,65 @@
     },
 }
 
-TEST_SIDECAR_DRIVER_APPLICATION = {
+TEST_DRIVER_WITH_SIDECAR_APPLICATION = {
     "apiVersion": "sparkoperator.k8s.io/v1beta2",
     "kind": "SparkApplication",
     "metadata": {
-        "creationTimestamp": "2020-02-27T08:03:02Z",
-        "generation": 4,
-        "name": "spark-pi",
+        "creationTimestamp": "2020-02-24T07:34:22Z",
+        "generation": 1,
+        "labels": {"spark_flow_name": "spark-pi"},
+        "name": "spark-pi-2020-02-24-1",
         "namespace": "default",
-        "resourceVersion": "552073",
+        "resourceVersion": "455577",
         "selfLink": "/apis/sparkoperator.k8s.io/v1beta2/namespaces/default/sparkapplications/spark-pi",
-        "uid": "0c93527d-4dd9-4006-b40a-1672872e8d6f",
+        "uid": "9f825516-6e1a-4af1-8967-b05661e8fb08",
     },
     "spec": {
-        "arguments": ["100000"],
         "driver": {
             "coreLimit": "1200m",
             "cores": 1,
-            "labels": {"version": "2.4.4"},
+            "labels": {"spark_flow_name": "spark-pi", "version": "2.4.4"},
             "memory": "512m",
             "serviceAccount": "default",
-            "sidecars": {
-                "name": "sidecar-container",
-                "image": "busybox",
-                "command": "[\"sh\",\"-c\",\"for i in {2..10}; do echo \"output: $i\"; done\"]",
-            },
+            "volumeMounts": [{"mountPath": "/tmp", "name": "test-volume"}],
+            "sidecars": [{"name": "sidecar1", "image": "hello-world:latest"}],
         },
         "executor": {
             "cores": 1,
-            "instances": 1,
-            "labels": {"version": "2.4.4"},
+            "instances": 3,
+            "labels": {"spark_flow_name": "spark-pi", "version": "2.4.4"},
             "memory": "512m",
+            "volumeMounts": [{"mountPath": "/tmp", "name": "test-volume"}],
         },
-        "image": "gcr.io/spark-operator/spark:v2.4.4-gcs-prometheus",
+        "image": "gcr.io/spark-operator/spark:v2.4.4",
         "imagePullPolicy": "Always",
         "mainApplicationFile": "local:///opt/spark/examples/jars/spark-examples_2.11-2.4.4.jar",
         "mainClass": "org.apache.spark.examples.SparkPi",
         "mode": "cluster",
-        "monitoring": {
-            "exposeDriverMetrics": True,
-            "exposeExecutorMetrics": True,
-            "prometheus": {
-                "jmxExporterJar": "/prometheus/jmx_prometheus_javaagent-0.11.0.jar",
-                "port": 8090,
-            },
-        },
         "restartPolicy": {"type": "Never"},
         "sparkVersion": "2.4.4",
         "type": "Scala",
+        "volumes": [{"hostPath": {"path": "/tmp", "type": "Directory"}, "name": "test-volume"}],
     },
     "status": {
-        "applicationState": {"state": "PENDING_RERUN"},
-        "driverInfo": {},
-        "lastSubmissionAttemptTime": None,
-        "terminationTime": None,
+        "applicationState": {"state": "COMPLETED"},
+        "driverInfo": {
+            "podName": "spark-pi-2020-02-24-1-driver",
+            "webUIAddress": "10.97.130.44:4040",
+            "webUIPort": 4040,
+            "webUIServiceName": "spark-pi-2020-02-24-1-ui-svc",
+        },
+        "executionAttempts": 1,
+        "executorState": {
+            "spark-pi-2020-02-24-1-1582529666227-exec-1": "FAILED",
+            "spark-pi-2020-02-24-1-1582529666227-exec-2": "FAILED",
+            "spark-pi-2020-02-24-1-1582529666227-exec-3": "FAILED",
+        },
+        "lastSubmissionAttemptTime": "2020-02-24T07:34:30Z",
+        "sparkApplicationId": "spark-7bb432c422ca46f3854838c419460fec",
+        "submissionAttempts": 1,
+        "submissionID": "1a1f9c5e-6bdd-4824-806f-40a814c1cf43",
+        "terminationTime": "2020-02-24T07:35:01Z",
     },
 }
 
@@ -783,8 +788,9 @@ def test_driver_logging_failure(
         )
         with pytest.raises(AirflowException):
             sensor.poke(None)
-        mock_log_call.assert_called_once_with("spark-pi-driver", namespace="default",
-                                              container='spark-kubernetes-driver')
+        mock_log_call.assert_called_once_with(
+            "spark-pi-driver", namespace="default", container='spark-kubernetes-driver'
+        )
         error_log_call.assert_called_once_with(TEST_POD_LOG_RESULT)
 
     @patch(
@@ -806,8 +812,9 @@ def test_driver_logging_completed(
             task_id="test_task_id",
         )
         sensor.poke(None)
-        mock_log_call.assert_called_once_with("spark-pi-2020-02-24-1-driver", namespace="default",
-                                              container='spark-kubernetes-driver')
+        mock_log_call.assert_called_once_with(
+            "spark-pi-2020-02-24-1-driver", namespace="default", container='spark-kubernetes-driver'
+        )
         log_info_call = info_log_call.mock_calls[2]
         log_value = log_info_call[1][0]
         assert log_value == TEST_POD_LOG_RESULT
@@ -835,7 +842,7 @@ def test_driver_logging_error(
 
     @patch(
         "kubernetes.client.api.custom_objects_api.CustomObjectsApi.get_namespaced_custom_object",
-        return_value=TEST_SIDECAR_DRIVER_APPLICATION,
+        return_value=TEST_DRIVER_WITH_SIDECAR_APPLICATION,
     )
     @patch("logging.Logger.info")
     @patch(
@@ -852,7 +859,9 @@ def test_sidecar_driver_logging_completed(
             task_id="test_task_id",
         )
         sensor.poke(None)
-        mock_log_call.assert_called_once_with("spark-pi-2020-02-24-1-driver", namespace="default")
+        mock_log_call.assert_called_once_with(
+            "spark-pi-2020-02-24-1-driver", namespace="default", container='spark-kubernetes-driver'
+        )
         log_info_call = info_log_call.mock_calls[2]
         log_value = log_info_call[1][0]
         assert log_value == TEST_POD_LOG_RESULT
