From afc47dacd5e544546bc6639ca34391b23041c40c Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sat, 13 Jul 2024 01:54:01 +0900
Subject: [PATCH 01/25] Add: added kerberos related connection
 fields(principal, keytab)

---
 .../apache/spark/hooks/spark_submit.py        | 71 ++++++++++++++++---
 1 file changed, 63 insertions(+), 8 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index ef0c4b5d9dd37..e8af27cf9a8cb 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -17,11 +17,16 @@
 # under the License.
 from __future__ import annotations
 
+import base64
 import contextlib
 import os
 import re
+import shutil
 import subprocess
+import tempfile
 import time
+import uuid
+from pathlib import Path
 from typing import Any, Iterator
 
 from airflow.configuration import conf as airflow_conf
@@ -67,7 +72,9 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
     :param executor_memory: Memory per executor (e.g. 1000M, 2G) (Default: 1G)
     :param driver_memory: Memory allocated to the driver (e.g. 1000M, 2G) (Default: 1G)
     :param keytab: Full path to the file that contains the keytab
+                        (will overwrite any keytab defined in the connection's extra JSON)
     :param principal: The name of the kerberos principal used for keytab
+                        (will overwrite any principal defined in the connection's extra JSON)
     :param proxy_user: User to impersonate when submitting the application
     :param name: Name of the job (default airflow-spark)
     :param num_executors: Number of executors to launch
@@ -101,14 +108,17 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]:
         return {
             "hidden_fields": ["schema", "login", "password", "extra"],
             "relabeling": {},
+            "placeholders": {
+                "keytab": "<base64 encoded Keytab Content>",
+            },
         }
 
     @classmethod
     def get_connection_form_widgets(cls) -> dict[str, Any]:
         """Return connection widgets to add to Spark connection form."""
-        from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
+        from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget
         from flask_babel import lazy_gettext
-        from wtforms import StringField
+        from wtforms import PasswordField, StringField
         from wtforms.validators import Optional, any_of
 
         return {
@@ -135,6 +145,17 @@ def get_connection_form_widgets(cls) -> dict[str, Any]:
             "namespace": StringField(
                 lazy_gettext("Kubernetes namespace"), widget=BS3TextFieldWidget(), validators=[Optional()]
             ),
+            "principal": StringField(
+                lazy_gettext("Principal"),
+                widget=BS3TextFieldWidget(),
+                validators=[Optional()],
+            ),
+            "keytab": PasswordField(
+                lazy_gettext("Keytab"),
+                widget=BS3PasswordFieldWidget(),
+                description="Run the command `base64 <your-keytab-path>` and use its output.",
+                validators=[Optional()],
+            ),
         }
 
     def __init__(
@@ -236,6 +257,8 @@ def _resolve_connection(self) -> dict[str, Any]:
             "deploy_mode": None,
             "spark_binary": self.spark_binary or DEFAULT_SPARK_BINARY,
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
 
         try:
@@ -268,6 +291,12 @@ def _resolve_connection(self) -> dict[str, Any]:
                 )
             conn_data["spark_binary"] = self.spark_binary
             conn_data["namespace"] = extra.get("namespace")
+            conn_data["principal"] = self._principal if self._principal else extra.get("principal")
+            base64_keytab = extra.get("keytab")
+            if self._keytab is not None:
+                conn_data["keytab"] = self._keytab
+            elif base64_keytab is not None:
+                conn_data["keytab"] = self._get_keytab_from_base64(base64_keytab)
         except AirflowException:
             self.log.info(
                 "Could not load connection string %s, defaulting to %s", self._conn_id, conn_data["master"]
@@ -281,6 +310,30 @@ def _resolve_connection(self) -> dict[str, Any]:
     def get_conn(self) -> Any:
         pass
 
+    def _get_keytab_from_base64(self, base64_keytab: str) -> str:
+        _uuid = uuid.uuid4()
+        temp_dir_path = Path(tempfile.gettempdir()).resolve()
+        temp_file_name = f"airflow_keytab-{self._connection.get('principal', _uuid)}"
+
+        keytab_path = temp_dir_path / temp_file_name
+        staging_path = temp_dir_path / f".{temp_file_name}.{_uuid}"
+
+        try:
+            with open(staging_path, "wb") as f:
+                self.log.info("Saving keytab to %s", staging_path)
+                f.write(base64.b64decode(base64_keytab))
+
+            self.log.info("Moving keytab from %s to %s", staging_path, keytab_path)
+            shutil.move(staging_path, keytab_path)
+            return str(keytab_path)
+        except Exception as err:
+            self.log.error("Failed to save keytab: %s", err)
+            raise
+        finally:
+            if staging_path.exists():
+                self.log.info("Removing staging keytab file: %s", staging_path)
+                staging_path.unlink()
+
     def _get_spark_binary_path(self) -> list[str]:
         # Assume that spark-submit is present in the path to the executing user
         return [self._connection["spark_binary"]]
@@ -369,10 +422,10 @@ def _build_spark_submit_command(self, application: str) -> list[str]:
             connection_cmd += ["--executor-memory", self._executor_memory]
         if self._driver_memory:
             connection_cmd += ["--driver-memory", self._driver_memory]
-        if self._keytab:
-            connection_cmd += ["--keytab", self._keytab]
-        if self._principal:
-            connection_cmd += ["--principal", self._principal]
+        if self._connection["keytab"]:
+            connection_cmd += ["--keytab", self._connection["keytab"]]
+        if self._connection["principal"]:
+            connection_cmd += ["--principal", self._connection["principal"]]
         if self._use_krb5ccache:
             if not os.getenv("KRB5CCNAME"):
                 raise AirflowException(
@@ -702,11 +755,13 @@ def on_kill(self) -> None:
             if self._yarn_application_id:
                 kill_cmd = f"yarn application -kill {self._yarn_application_id}".split()
                 env = {**os.environ, **(self._env or {})}
-                if self._keytab is not None and self._principal is not None:
+                if self._connection["keytab"] is not None and self._connection["principal"] is not None:
                     # we are ignoring renewal failures from renew_from_kt
                     # here as the failure could just be due to a non-renewable ticket,
                     # we still attempt to kill the yarn application
-                    renew_from_kt(self._principal, self._keytab, exit_on_fail=False)
+                    renew_from_kt(
+                        self._connection["principal"], self._connection["keytab"], exit_on_fail=False
+                    )
                     env = os.environ.copy()
                     ccacche = airflow_conf.get_mandatory_value("kerberos", "ccache")
                     env["KRB5CCNAME"] = ccacche

From 4e7569aa4ae20c99019dfc9fe90bc9efc71b86a0 Mon Sep 17 00:00:00 2001
From: "Youngha, Park" <32288527+softyoungha@users.noreply.github.com>
Date: Sat, 13 Jul 2024 02:14:47 +0900
Subject: [PATCH 02/25] Update
 airflow/providers/apache/spark/hooks/spark_submit.py

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index e8af27cf9a8cb..524c8ee2d7726 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -292,10 +292,9 @@ def _resolve_connection(self) -> dict[str, Any]:
             conn_data["spark_binary"] = self.spark_binary
             conn_data["namespace"] = extra.get("namespace")
             conn_data["principal"] = self._principal if self._principal else extra.get("principal")
+            conn_data["keytab"] = self._keytab
             base64_keytab = extra.get("keytab")
-            if self._keytab is not None:
-                conn_data["keytab"] = self._keytab
-            elif base64_keytab is not None:
+            if base64_keytab is not None:
                 conn_data["keytab"] = self._get_keytab_from_base64(base64_keytab)
         except AirflowException:
             self.log.info(

From 1ec7a1cb7f95204d67555887b6aa4538424f0521 Mon Sep 17 00:00:00 2001
From: "Youngha, Park" <32288527+softyoungha@users.noreply.github.com>
Date: Sat, 13 Jul 2024 02:15:07 +0900
Subject: [PATCH 03/25] Update
 airflow/providers/apache/spark/hooks/spark_submit.py

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 524c8ee2d7726..4144df69aa9f6 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -291,7 +291,7 @@ def _resolve_connection(self) -> dict[str, Any]:
                 )
             conn_data["spark_binary"] = self.spark_binary
             conn_data["namespace"] = extra.get("namespace")
-            conn_data["principal"] = self._principal if self._principal else extra.get("principal")
+            conn_data["principal"] = self._principal or extra.get("principal")
             conn_data["keytab"] = self._keytab
             base64_keytab = extra.get("keytab")
             if base64_keytab is not None:

From afb28cb49de730bda5bd9d422e61ecd67ade3e43 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sat, 13 Jul 2024 02:35:49 +0900
Subject: [PATCH 04/25] fixed wrong logic since keytab priority

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 4144df69aa9f6..5ffda202143be 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -292,9 +292,10 @@ def _resolve_connection(self) -> dict[str, Any]:
             conn_data["spark_binary"] = self.spark_binary
             conn_data["namespace"] = extra.get("namespace")
             conn_data["principal"] = self._principal or extra.get("principal")
-            conn_data["keytab"] = self._keytab
             base64_keytab = extra.get("keytab")
-            if base64_keytab is not None:
+            if self._keytab is not None:
+                conn_data["keytab"] = self._keytab
+            elif base64_keytab is not None:
                 conn_data["keytab"] = self._get_keytab_from_base64(base64_keytab)
         except AirflowException:
             self.log.info(

From 0a9d7d07f3e01816a3160c942b44e1c470a89019 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sat, 13 Jul 2024 06:51:46 +0900
Subject: [PATCH 05/25] fixed wrong temp_file_name logic

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 5ffda202143be..07852fe869e2e 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -313,7 +313,7 @@ def get_conn(self) -> Any:
     def _get_keytab_from_base64(self, base64_keytab: str) -> str:
         _uuid = uuid.uuid4()
         temp_dir_path = Path(tempfile.gettempdir()).resolve()
-        temp_file_name = f"airflow_keytab-{self._connection.get('principal', _uuid)}"
+        temp_file_name = f"airflow_keytab-{self._connection['principal'] or _uuid}"
 
         keytab_path = temp_dir_path / temp_file_name
         staging_path = temp_dir_path / f".{temp_file_name}.{_uuid}"

From feb3398ee764a30e745e76c0b1e90f506b1fca58 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sat, 13 Jul 2024 06:58:10 +0900
Subject: [PATCH 06/25] Added annotations on operators/spark_submit.py

---
 airflow/providers/apache/spark/operators/spark_submit.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/airflow/providers/apache/spark/operators/spark_submit.py b/airflow/providers/apache/spark/operators/spark_submit.py
index 281919b2b14c8..16f232dd02946 100644
--- a/airflow/providers/apache/spark/operators/spark_submit.py
+++ b/airflow/providers/apache/spark/operators/spark_submit.py
@@ -58,7 +58,9 @@ class SparkSubmitOperator(BaseOperator):
     :param executor_memory: Memory per executor (e.g. 1000M, 2G) (Default: 1G)
     :param driver_memory: Memory allocated to the driver (e.g. 1000M, 2G) (Default: 1G)
     :param keytab: Full path to the file that contains the keytab (templated)
+                        (will overwrite any keytab defined in the connection's extra JSON)
     :param principal: The name of the kerberos principal used for keytab (templated)
+                        (will overwrite any principal defined in the connection's extra JSON)
     :param proxy_user: User to impersonate when submitting the application (templated)
     :param name: Name of the job (default airflow-spark). (templated)
     :param num_executors: Number of executors to launch

From 8148d5c11f99fc11145eb97ed253424b02b10335 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sat, 13 Jul 2024 07:26:19 +0900
Subject: [PATCH 07/25] changed raise Exception to AirflowException

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 07852fe869e2e..643fe4b24711a 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -328,7 +328,7 @@ def _get_keytab_from_base64(self, base64_keytab: str) -> str:
             return str(keytab_path)
         except Exception as err:
             self.log.error("Failed to save keytab: %s", err)
-            raise
+            raise AirflowException("Failed to save keytab") from err
         finally:
             if staging_path.exists():
                 self.log.info("Removing staging keytab file: %s", staging_path)

From 30606bc3c292fb1177a25c518ae46023f5f4666d Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 05:12:41 +0900
Subject: [PATCH 08/25] Fixed principal referencing bug

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 643fe4b24711a..68ebcd9a8b074 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -296,7 +296,7 @@ def _resolve_connection(self) -> dict[str, Any]:
             if self._keytab is not None:
                 conn_data["keytab"] = self._keytab
             elif base64_keytab is not None:
-                conn_data["keytab"] = self._get_keytab_from_base64(base64_keytab)
+                conn_data["keytab"] = self._get_keytab_from_base64(base64_keytab, conn_data["principal"])
         except AirflowException:
             self.log.info(
                 "Could not load connection string %s, defaulting to %s", self._conn_id, conn_data["master"]
@@ -310,10 +310,10 @@ def _resolve_connection(self) -> dict[str, Any]:
     def get_conn(self) -> Any:
         pass
 
-    def _get_keytab_from_base64(self, base64_keytab: str) -> str:
+    def _get_keytab_from_base64(self, base64_keytab: str, principal: str | None) -> str:
         _uuid = uuid.uuid4()
         temp_dir_path = Path(tempfile.gettempdir()).resolve()
-        temp_file_name = f"airflow_keytab-{self._connection['principal'] or _uuid}"
+        temp_file_name = f"airflow_keytab-{principal or _uuid}"
 
         keytab_path = temp_dir_path / temp_file_name
         staging_path = temp_dir_path / f".{temp_file_name}.{_uuid}"

From 6d37442c3e82e296e999a8f2db5fca7e7213c25f Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 05:50:35 +0900
Subject: [PATCH 09/25] Added checksum validating logic to do not rewrite
 keytab

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 68ebcd9a8b074..463cb1ad11bfe 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -19,6 +19,7 @@
 
 import base64
 import contextlib
+import hashlib
 import os
 import re
 import shutil
@@ -318,6 +319,16 @@ def _get_keytab_from_base64(self, base64_keytab: str, principal: str | None) ->
         keytab_path = temp_dir_path / temp_file_name
         staging_path = temp_dir_path / f".{temp_file_name}.{_uuid}"
 
+        # validate exists keytab file
+        if keytab_path.exists():
+            with open(keytab_path, "rb") as f:
+                old_checksum = hashlib.sha256(f.read()).hexdigest()
+            new_checksum = hashlib.sha256(base64.b64decode(base64_keytab)).hexdigest()
+            if old_checksum == new_checksum:
+                self.log.info("Keytab file already exists and is the same as the provided keytab")
+                return str(keytab_path)
+
+        # write new keytab file
         try:
             with open(staging_path, "wb") as f:
                 self.log.info("Saving keytab to %s", staging_path)

From 6636d528d2934afcce3ca1d32a4535ea51cf3049 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 05:58:28 +0900
Subject: [PATCH 10/25] Fixed arg-type error

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 463cb1ad11bfe..dd451094dc799 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -319,11 +319,13 @@ def _get_keytab_from_base64(self, base64_keytab: str, principal: str | None) ->
         keytab_path = temp_dir_path / temp_file_name
         staging_path = temp_dir_path / f".{temp_file_name}.{_uuid}"
 
+        keytab = base64.b64decode(base64_keytab)
+
         # validate exists keytab file
         if keytab_path.exists():
             with open(keytab_path, "rb") as f:
                 old_checksum = hashlib.sha256(f.read()).hexdigest()
-            new_checksum = hashlib.sha256(base64.b64decode(base64_keytab)).hexdigest()
+            new_checksum = hashlib.sha256(keytab).hexdigest()
             if old_checksum == new_checksum:
                 self.log.info("Keytab file already exists and is the same as the provided keytab")
                 return str(keytab_path)
@@ -332,10 +334,10 @@ def _get_keytab_from_base64(self, base64_keytab: str, principal: str | None) ->
         try:
             with open(staging_path, "wb") as f:
                 self.log.info("Saving keytab to %s", staging_path)
-                f.write(base64.b64decode(base64_keytab))
+                f.write(keytab)
 
             self.log.info("Moving keytab from %s to %s", staging_path, keytab_path)
-            shutil.move(staging_path, keytab_path)
+            shutil.move(str(staging_path), str(keytab_path))
             return str(keytab_path)
         except Exception as err:
             self.log.error("Failed to save keytab: %s", err)

From 5947c96a91f983898124a875fae6958159d6449e Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 07:22:25 +0900
Subject: [PATCH 11/25] simplified existing keytab compare logic

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index dd451094dc799..fa05032c01936 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -19,7 +19,6 @@
 
 import base64
 import contextlib
-import hashlib
 import os
 import re
 import shutil
@@ -324,9 +323,8 @@ def _get_keytab_from_base64(self, base64_keytab: str, principal: str | None) ->
         # validate exists keytab file
         if keytab_path.exists():
             with open(keytab_path, "rb") as f:
-                old_checksum = hashlib.sha256(f.read()).hexdigest()
-            new_checksum = hashlib.sha256(keytab).hexdigest()
-            if old_checksum == new_checksum:
+                existing_keytab = f.read()
+            if existing_keytab == keytab:
                 self.log.info("Keytab file already exists and is the same as the provided keytab")
                 return str(keytab_path)
 

From 51032ed3ffeeb0349a84bca18f548623f11a3c6c Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 07:23:00 +0900
Subject: [PATCH 12/25] Added principal test cases

---
 .../apache/spark/hooks/test_spark_submit.py   | 52 +++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py
index fbd2c9999a754..bfd74301c8ee7 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -143,6 +143,14 @@ def setup_method(self):
                 extra='{"deploy-mode": "client"}',
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id="spark_principal_set",
+                conn_type="spark",
+                host="yarn",
+                extra='{"principal": "user/spark@airflow.org"}',
+            )
+        )
 
     @patch(
         "airflow.providers.apache.spark.hooks.spark_submit.os.getenv", return_value="/tmp/airflow_krb5_ccache"
@@ -554,6 +562,50 @@ def test_resolve_connection_spark_standalone_cluster_connection(self):
         assert connection == expected_spark_connection
         assert cmd[0] == "spark-submit"
 
+    def test_resolve_connection_principal_set_connection(self):
+        # Given
+        hook = SparkSubmitHook(conn_id="spark_principal_set")
+
+        # When
+        connection = hook._resolve_connection()
+        cmd = hook._build_spark_submit_command(self._spark_job_file)
+
+        # Then
+        dict_cmd = self.cmd_args_to_dict(cmd)
+        expected_spark_connection = {
+            "master": "yarn",
+            "spark_binary": "spark-submit",
+            "deploy_mode": None,
+            "queue": None,
+            "namespace": None,
+            "principal": "user/spark@airflow.org",
+            "keytab": None,
+        }
+        assert connection == expected_spark_connection
+        assert dict_cmd["--principal"] == "user/spark@airflow.org"
+
+    def test_resolve_connection_principal_value_override(self):
+        # Given
+        hook = SparkSubmitHook(conn_id="spark_principal_set", principal="will-override")
+
+        # When
+        connection = hook._resolve_connection()
+        cmd = hook._build_spark_submit_command(self._spark_job_file)
+
+        # Then
+        dict_cmd = self.cmd_args_to_dict(cmd)
+        expected_spark_connection = {
+            "master": "yarn",
+            "spark_binary": "spark-submit",
+            "deploy_mode": None,
+            "queue": None,
+            "namespace": None,
+            "principal": "will-override",
+            "keytab": None,
+        }
+        assert connection == expected_spark_connection
+        assert dict_cmd["--principal"] == "will-override"
+
     def test_resolve_spark_submit_env_vars_standalone_client_mode(self):
         # Given
         hook = SparkSubmitHook(conn_id="spark_standalone_cluster_client_mode", env_vars={"bar": "foo"})

From 3c0e60adacebec06c74a7e2b9f78346c11b2dbdc Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 09:11:17 +0900
Subject: [PATCH 13/25] Added keytab test cases

---
 .../apache/spark/hooks/test_spark_submit.py   | 147 +++++++++++++++++-
 1 file changed, 146 insertions(+), 1 deletion(-)

diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py
index ba365cbc8c89c..6951e1b5ce38c 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -17,9 +17,11 @@
 # under the License.
 from __future__ import annotations
 
+import base64
 import os
 from io import StringIO
-from unittest.mock import call, patch
+from pathlib import Path
+from unittest.mock import call, mock_open, patch
 
 import pytest
 
@@ -151,6 +153,14 @@ def setup_method(self):
                 extra='{"principal": "user/spark@airflow.org"}',
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id="spark_keytab_set",
+                conn_type="spark",
+                host="yarn",
+                extra='{"keytab": "privileged_user.keytab"}',
+            )
+        )
 
     @patch(
         "airflow.providers.apache.spark.hooks.spark_submit.os.getenv", return_value="/tmp/airflow_krb5_ccache"
@@ -606,6 +616,54 @@ def test_resolve_connection_principal_value_override(self):
         assert connection == expected_spark_connection
         assert dict_cmd["--principal"] == "will-override"
 
+    @patch(
+        "airflow.providers.apache.spark.hooks.spark_submit.SparkSubmitHook._get_keytab_from_base64",
+        return_value="privileged_user.keytab",
+    )
+    def test_resolve_connection_keytab_set_connection(self, mock_get_keytab_from_base64):
+        # Given
+        hook = SparkSubmitHook(conn_id="spark_keytab_set")
+
+        # When
+        connection = hook._resolve_connection()
+        cmd = hook._build_spark_submit_command(self._spark_job_file)
+
+        # Then
+        dict_cmd = self.cmd_args_to_dict(cmd)
+        expected_spark_connection = {
+            "master": "yarn",
+            "spark_binary": "spark-submit",
+            "deploy_mode": None,
+            "queue": None,
+            "namespace": None,
+            "principal": None,
+            "keytab": "privileged_user.keytab",
+        }
+        assert connection == expected_spark_connection
+        assert dict_cmd["--keytab"] == "privileged_user.keytab"
+
+    def test_resolve_connection_keytab_value_override(self):
+        # Given
+        hook = SparkSubmitHook(conn_id="spark_keytab_set", keytab="will-override")
+
+        # When
+        connection = hook._resolve_connection()
+        cmd = hook._build_spark_submit_command(self._spark_job_file)
+
+        # Then
+        dict_cmd = self.cmd_args_to_dict(cmd)
+        expected_spark_connection = {
+            "master": "yarn",
+            "spark_binary": "spark-submit",
+            "deploy_mode": None,
+            "queue": None,
+            "namespace": None,
+            "principal": None,
+            "keytab": "will-override",
+        }
+        assert connection == expected_spark_connection
+        assert dict_cmd["--keytab"] == "will-override"
+
     def test_resolve_spark_submit_env_vars_standalone_client_mode(self):
         # Given
         hook = SparkSubmitHook(conn_id="spark_standalone_cluster_client_mode", env_vars={"bar": "foo"})
@@ -945,3 +1003,90 @@ def test_masks_passwords(self, command: str, expected: str) -> None:
 
         # Then
         assert command_masked == expected
+
+    @patch("airflow.providers.apache.spark.hooks.spark_submit.uuid.uuid4")
+    @patch("pathlib.Path.resolve")
+    @patch("airflow.providers.apache.spark.hooks.spark_submit.shutil.move")
+    @patch("pathlib.Path.exists")
+    @patch("builtins.open", new_callable=mock_open)
+    def test_get_keytab_from_base64_with_new_keytab(
+        self,
+        mock_open,
+        mock_exists,
+        mock_move,
+        mock_resolve,
+        mock_uuid4,
+    ):
+        # Given
+        hook = SparkSubmitHook()
+
+        keytab_value = b"abcd"
+        base64_keytab = base64.b64encode(keytab_value).decode("UTF-8")
+        mock_uuid4.return_value = "uuid"
+        mock_resolve.return_value = Path("resolved_path")
+        mock_exists.return_value = False
+
+        # When
+        keytab = hook._get_keytab_from_base64(base64_keytab, None)
+
+        # Then
+        assert keytab == "resolved_path/airflow_keytab-uuid"
+        mock_open().write.assert_called_once_with(keytab_value)
+        mock_move.assert_called_once()
+
+    @patch("pathlib.Path.resolve")
+    @patch("airflow.providers.apache.spark.hooks.spark_submit.shutil.move")
+    @patch("pathlib.Path.exists")
+    @patch("builtins.open", new_callable=mock_open)
+    def test_get_keytab_from_base64_with_new_keytab_with_principal(
+        self,
+        mock_open,
+        mock_exists,
+        mock_move,
+        mock_resolve,
+    ):
+        # Given
+        hook = SparkSubmitHook()
+
+        principal = "user/spark@airflow.org"
+        keytab_value = b"abcd"
+        base64_keytab = base64.b64encode(keytab_value).decode("UTF-8")
+        mock_resolve.return_value = Path("resolved_path")
+        mock_exists.return_value = False
+
+        # When
+        keytab = hook._get_keytab_from_base64(base64_keytab, principal)
+
+        # Then
+        assert keytab == f"resolved_path/airflow_keytab-{principal}"
+        mock_open().write.assert_called_once_with(keytab_value)
+        mock_move.assert_called_once()
+
+    @patch("pathlib.Path.resolve")
+    @patch("pathlib.Path.exists")
+    @patch("builtins.open", new_callable=mock_open)
+    def test_get_keytab_from_base64_with_existing_keytab(
+        self,
+        mock_open,
+        mock_exists,
+        mock_resolve,
+    ):
+        # Given
+        hook = SparkSubmitHook()
+
+        principal = "user/spark@airflow.org"
+        keytab_value = b"abcd"
+        base64_keytab = base64.b64encode(keytab_value)
+        mock_resolve.return_value = Path("resolved_path")
+        mock_exists.return_value = True
+        _mock_open = mock_open()
+        _mock_open.read.return_value = base64.b64decode(base64_keytab)
+
+        # When
+        keytab = hook._get_keytab_from_base64(base64_keytab.decode("UTF-8"), principal)
+
+        # Then
+        assert keytab == f"resolved_path/airflow_keytab-{principal}"
+        mock_open.assert_called_with(Path(f"resolved_path/airflow_keytab-{principal}"), "rb")
+        _mock_open.read.assert_called_once()
+        assert not _mock_open.write.called, "Keytab file should not be written"

From e21e4f494a68c79093a1a779662ab82589095e50 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 09:22:57 +0900
Subject: [PATCH 14/25] added principal & keytab to expected conn values

---
 .../apache/spark/hooks/test_spark_submit.py   | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py
index 6951e1b5ce38c..510de537866db 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -353,6 +353,8 @@ def test_resolve_connection_yarn_default(self):
             "deploy_mode": None,
             "queue": None,
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert dict_cmd["--master"] == "yarn"
@@ -373,6 +375,8 @@ def test_resolve_connection_yarn_default_connection(self):
             "deploy_mode": None,
             "queue": "root.default",
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert dict_cmd["--master"] == "yarn"
@@ -394,6 +398,8 @@ def test_resolve_connection_mesos_default_connection(self):
             "deploy_mode": None,
             "queue": None,
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert dict_cmd["--master"] == "mesos://host:5050"
@@ -414,6 +420,8 @@ def test_resolve_connection_spark_yarn_cluster_connection(self):
             "deploy_mode": "cluster",
             "queue": "root.etl",
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert dict_cmd["--master"] == "yarn://yarn-master"
@@ -436,6 +444,8 @@ def test_resolve_connection_spark_k8s_cluster_connection(self):
             "master": "k8s://https://k8s-master",
             "deploy_mode": "cluster",
             "namespace": "mynamespace",
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert dict_cmd["--master"] == "k8s://https://k8s-master"
@@ -460,6 +470,8 @@ def test_resolve_connection_spark_k8s_cluster_ns_conf(self):
             "master": "k8s://https://k8s-master",
             "deploy_mode": "cluster",
             "namespace": "airflow",
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert dict_cmd["--master"] == "k8s://https://k8s-master"
@@ -481,6 +493,8 @@ def test_resolve_connection_spark_binary_set_connection(self):
             "deploy_mode": None,
             "queue": None,
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert cmd[0] == "spark2-submit"
@@ -500,6 +514,8 @@ def test_resolve_connection_spark_binary_spark3_submit_set_connection(self):
             "deploy_mode": None,
             "queue": None,
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert cmd[0] == "spark3-submit"
@@ -530,6 +546,8 @@ def test_resolve_connection_spark_binary_default_value_override(self):
             "deploy_mode": None,
             "queue": None,
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert cmd[0] == "spark3-submit"
@@ -549,6 +567,8 @@ def test_resolve_connection_spark_binary_default_value(self):
             "deploy_mode": None,
             "queue": "root.default",
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert cmd[0] == "spark-submit"
@@ -568,6 +588,8 @@ def test_resolve_connection_spark_standalone_cluster_connection(self):
             "deploy_mode": "cluster",
             "queue": None,
             "namespace": None,
+            "principal": None,
+            "keytab": None,
         }
         assert connection == expected_spark_connection
         assert cmd[0] == "spark-submit"

From 24a1c78e57bd9ded3d568f85c9da18861e34a8d7 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 09:51:50 +0900
Subject: [PATCH 15/25] added principal & keytab getter for the borken
 conneciton case

---
 .../apache/spark/hooks/spark_submit.py        | 26 +++++++++++++------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 0ddd3e81f9482..cc460db530917 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -349,6 +349,14 @@ def _get_spark_binary_path(self) -> list[str]:
         # Assume that spark-submit is present in the path to the executing user
         return [self._connection["spark_binary"]]
 
+    def _get_principal(self) -> str | None:
+        # for the case where the principal is not set in the broken connection
+        return self._connection["principal"] or self._principal
+
+    def _get_keytab(self) -> str | None:
+        # for the case where the keytab is not set in the broken connection
+        return self._connection["keytab"] or self._keytab
+
     def _mask_cmd(self, connection_cmd: str | list[str]) -> str:
         # Mask any password related fields in application args with key value pair
         # where key contains password (case insensitive), e.g. HivePassword='abc'
@@ -433,10 +441,12 @@ def _build_spark_submit_command(self, application: str) -> list[str]:
             connection_cmd += ["--executor-memory", self._executor_memory]
         if self._driver_memory:
             connection_cmd += ["--driver-memory", self._driver_memory]
-        if self._connection["keytab"]:
-            connection_cmd += ["--keytab", self._connection["keytab"]]
-        if self._connection["principal"]:
-            connection_cmd += ["--principal", self._connection["principal"]]
+        keytab = self._get_keytab()
+        if keytab:
+            connection_cmd += ["--keytab", keytab]
+        principal = self._get_principal()
+        if principal:
+            connection_cmd += ["--principal", principal]
         if self._use_krb5ccache:
             if not os.getenv("KRB5CCNAME"):
                 raise AirflowException(
@@ -766,13 +776,13 @@ def on_kill(self) -> None:
             if self._yarn_application_id:
                 kill_cmd = f"yarn application -kill {self._yarn_application_id}".split()
                 env = {**os.environ, **(self._env or {})}
-                if self._connection["keytab"] is not None and self._connection["principal"] is not None:
+                principal = self._get_principal()
+                keytab = self._get_keytab()
+                if keytab is not None and principal is not None:
                     # we are ignoring renewal failures from renew_from_kt
                     # here as the failure could just be due to a non-renewable ticket,
                     # we still attempt to kill the yarn application
-                    renew_from_kt(
-                        self._connection["principal"], self._connection["keytab"], exit_on_fail=False
-                    )
+                    renew_from_kt(principal, keytab, exit_on_fail=False)
                     env = os.environ.copy()
                     ccacche = airflow_conf.get_mandatory_value("kerberos", "ccache")
                     env["KRB5CCNAME"] = ccacche

From 872dbbe284f0139fb4101701fd15288b9b19302d Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sun, 14 Jul 2024 12:23:43 +0900
Subject: [PATCH 16/25] added guard logic for the case that failed to b64decode

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index cc460db530917..e74b4beda621f 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -318,7 +318,11 @@ def _get_keytab_from_base64(self, base64_keytab: str, principal: str | None) ->
         keytab_path = temp_dir_path / temp_file_name
         staging_path = temp_dir_path / f".{temp_file_name}.{_uuid}"
 
-        keytab = base64.b64decode(base64_keytab)
+        try:
+            keytab = base64.b64decode(base64_keytab)
+        except Exception as err:
+            self.log.error("Failed to decode base64 keytab: %s", err)
+            raise AirflowException("Failed to decode base64 keytab") from err
 
         # validate exists keytab file
         if keytab_path.exists():

From 11ae140746c50c77d006e0992c609fa19e643645 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Mon, 15 Jul 2024 02:16:34 +0900
Subject: [PATCH 17/25] removed useless logic on test code

---
 tests/providers/apache/spark/hooks/test_spark_submit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py
index 510de537866db..507f1cc02e018 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -1102,7 +1102,7 @@ def test_get_keytab_from_base64_with_existing_keytab(
         mock_resolve.return_value = Path("resolved_path")
         mock_exists.return_value = True
         _mock_open = mock_open()
-        _mock_open.read.return_value = base64.b64decode(base64_keytab)
+        _mock_open.read.return_value = keytab_value
 
         # When
         keytab = hook._get_keytab_from_base64(base64_keytab.decode("UTF-8"), principal)

From a9d4f1f946a48523f8f4c79817d3a3171202ddb5 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Mon, 15 Jul 2024 21:20:29 +0900
Subject: [PATCH 18/25] Mod: improved keytab_value_override test case

---
 tests/providers/apache/spark/hooks/test_spark_submit.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py
index 507f1cc02e018..072305bf93245 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -664,7 +664,8 @@ def test_resolve_connection_keytab_set_connection(self, mock_get_keytab_from_bas
         assert connection == expected_spark_connection
         assert dict_cmd["--keytab"] == "privileged_user.keytab"
 
-    def test_resolve_connection_keytab_value_override(self):
+    @patch("airflow.providers.apache.spark.hooks.spark_submit.SparkSubmitHook._get_keytab_from_base64")
+    def test_resolve_connection_keytab_value_override(self, mock_get_keytab_from_base64):
         # Given
         hook = SparkSubmitHook(conn_id="spark_keytab_set", keytab="will-override")
 
@@ -685,6 +686,7 @@ def test_resolve_connection_keytab_value_override(self):
         }
         assert connection == expected_spark_connection
         assert dict_cmd["--keytab"] == "will-override"
+        assert not mock_get_keytab_from_base64.called, "Should not call _get_keytab_from_base64"
 
     def test_resolve_spark_submit_env_vars_standalone_client_mode(self):
         # Given

From fbf520b88b63aca8a89fc6c1c308eb659f4acc59 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Tue, 16 Jul 2024 13:19:31 +0900
Subject: [PATCH 19/25] changed func name _get_keytab_from_base64 to
 _create_keytab_path_from_base64_keytab

---
 .../apache/spark/hooks/spark_submit.py        |  6 +++--
 .../apache/spark/hooks/test_spark_submit.py   | 26 +++++++++++--------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index e74b4beda621f..528a54f6651cd 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -296,7 +296,9 @@ def _resolve_connection(self) -> dict[str, Any]:
             if self._keytab is not None:
                 conn_data["keytab"] = self._keytab
             elif base64_keytab is not None:
-                conn_data["keytab"] = self._get_keytab_from_base64(base64_keytab, conn_data["principal"])
+                conn_data["keytab"] = self._create_keytab_path_from_base64_keytab(
+                    base64_keytab, conn_data["principal"]
+                )
         except AirflowException:
             self.log.info(
                 "Could not load connection string %s, defaulting to %s", self._conn_id, conn_data["master"]
@@ -310,7 +312,7 @@ def _resolve_connection(self) -> dict[str, Any]:
     def get_conn(self) -> Any:
         pass
 
-    def _get_keytab_from_base64(self, base64_keytab: str, principal: str | None) -> str:
+    def _create_keytab_path_from_base64_keytab(self, base64_keytab: str, principal: str | None) -> str:
         _uuid = uuid.uuid4()
         temp_dir_path = Path(tempfile.gettempdir()).resolve()
         temp_file_name = f"airflow_keytab-{principal or _uuid}"
diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py
index 072305bf93245..442597affab5c 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -639,10 +639,10 @@ def test_resolve_connection_principal_value_override(self):
         assert dict_cmd["--principal"] == "will-override"
 
     @patch(
-        "airflow.providers.apache.spark.hooks.spark_submit.SparkSubmitHook._get_keytab_from_base64",
+        "airflow.providers.apache.spark.hooks.spark_submit.SparkSubmitHook._create_keytab_path_from_base64_keytab",
         return_value="privileged_user.keytab",
     )
-    def test_resolve_connection_keytab_set_connection(self, mock_get_keytab_from_base64):
+    def test_resolve_connection_keytab_set_connection(self, mock_create_keytab_path_from_base64_keytab):
         # Given
         hook = SparkSubmitHook(conn_id="spark_keytab_set")
 
@@ -664,8 +664,10 @@ def test_resolve_connection_keytab_set_connection(self, mock_get_keytab_from_bas
         assert connection == expected_spark_connection
         assert dict_cmd["--keytab"] == "privileged_user.keytab"
 
-    @patch("airflow.providers.apache.spark.hooks.spark_submit.SparkSubmitHook._get_keytab_from_base64")
-    def test_resolve_connection_keytab_value_override(self, mock_get_keytab_from_base64):
+    @patch(
+        "airflow.providers.apache.spark.hooks.spark_submit.SparkSubmitHook._create_keytab_path_from_base64_keytab"
+    )
+    def test_resolve_connection_keytab_value_override(self, mock_create_keytab_path_from_base64_keytab):
         # Given
         hook = SparkSubmitHook(conn_id="spark_keytab_set", keytab="will-override")
 
@@ -686,7 +688,9 @@ def test_resolve_connection_keytab_value_override(self, mock_get_keytab_from_bas
         }
         assert connection == expected_spark_connection
         assert dict_cmd["--keytab"] == "will-override"
-        assert not mock_get_keytab_from_base64.called, "Should not call _get_keytab_from_base64"
+        assert (
+            not mock_create_keytab_path_from_base64_keytab.called
+        ), "Should not call _create_keytab_path_from_base64_keytab"
 
     def test_resolve_spark_submit_env_vars_standalone_client_mode(self):
         # Given
@@ -1033,7 +1037,7 @@ def test_masks_passwords(self, command: str, expected: str) -> None:
     @patch("airflow.providers.apache.spark.hooks.spark_submit.shutil.move")
     @patch("pathlib.Path.exists")
     @patch("builtins.open", new_callable=mock_open)
-    def test_get_keytab_from_base64_with_new_keytab(
+    def test_create_keytab_path_from_base64_keytab_with_new_keytab(
         self,
         mock_open,
         mock_exists,
@@ -1051,7 +1055,7 @@ def test_get_keytab_from_base64_with_new_keytab(
         mock_exists.return_value = False
 
         # When
-        keytab = hook._get_keytab_from_base64(base64_keytab, None)
+        keytab = hook._create_keytab_path_from_base64_keytab(base64_keytab, None)
 
         # Then
         assert keytab == "resolved_path/airflow_keytab-uuid"
@@ -1062,7 +1066,7 @@ def test_get_keytab_from_base64_with_new_keytab(
     @patch("airflow.providers.apache.spark.hooks.spark_submit.shutil.move")
     @patch("pathlib.Path.exists")
     @patch("builtins.open", new_callable=mock_open)
-    def test_get_keytab_from_base64_with_new_keytab_with_principal(
+    def test_create_keytab_path_from_base64_keytab_with_new_keytab_with_principal(
         self,
         mock_open,
         mock_exists,
@@ -1079,7 +1083,7 @@ def test_get_keytab_from_base64_with_new_keytab_with_principal(
         mock_exists.return_value = False
 
         # When
-        keytab = hook._get_keytab_from_base64(base64_keytab, principal)
+        keytab = hook._create_keytab_path_from_base64_keytab(base64_keytab, principal)
 
         # Then
         assert keytab == f"resolved_path/airflow_keytab-{principal}"
@@ -1089,7 +1093,7 @@ def test_get_keytab_from_base64_with_new_keytab_with_principal(
     @patch("pathlib.Path.resolve")
     @patch("pathlib.Path.exists")
     @patch("builtins.open", new_callable=mock_open)
-    def test_get_keytab_from_base64_with_existing_keytab(
+    def test_create_keytab_path_from_base64_keytab_with_existing_keytab(
         self,
         mock_open,
         mock_exists,
@@ -1107,7 +1111,7 @@ def test_get_keytab_from_base64_with_existing_keytab(
         _mock_open.read.return_value = keytab_value
 
         # When
-        keytab = hook._get_keytab_from_base64(base64_keytab.decode("UTF-8"), principal)
+        keytab = hook._create_keytab_path_from_base64_keytab(base64_keytab.decode("UTF-8"), principal)
 
         # Then
         assert keytab == f"resolved_path/airflow_keytab-{principal}"

From bcd3a3755c3c1f0349f03d9bf41f887c7dd0c4e7 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Mon, 2 Sep 2024 10:00:49 +0900
Subject: [PATCH 20/25] added additional comment on create_keytab parts

---
 airflow/providers/apache/spark/hooks/spark_submit.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index 1a92ca4040ba6..dc2bce3e36ece 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -327,7 +327,7 @@ def _create_keytab_path_from_base64_keytab(self, base64_keytab: str, principal:
             self.log.error("Failed to decode base64 keytab: %s", err)
             raise AirflowException("Failed to decode base64 keytab") from err
 
-        # validate exists keytab file
+        # If a keytab file with the same content exists, return its path for the --keytab argument.
         if keytab_path.exists():
             with open(keytab_path, "rb") as f:
                 existing_keytab = f.read()
@@ -335,7 +335,7 @@ def _create_keytab_path_from_base64_keytab(self, base64_keytab: str, principal:
                 self.log.info("Keytab file already exists and is the same as the provided keytab")
                 return str(keytab_path)
 
-        # write new keytab file
+        # Generate a new keytab file from the base64 encoded keytab value to use as a --keytab argument.
         try:
             with open(staging_path, "wb") as f:
                 self.log.info("Saving keytab to %s", staging_path)

From 912cab06280d8b6fa42279764cc28d9a5bda79ca Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Tue, 3 Sep 2024 17:11:53 +0900
Subject: [PATCH 21/25] added base64 decode exception test case

---
 tests/providers/apache/spark/hooks/test_spark_submit.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py
index 8046041f78177..ce2dcdfe6de42 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -1033,6 +1033,13 @@ def test_masks_passwords(self, command: str, expected: str) -> None:
         # Then
         assert command_masked == expected
 
+    def test_create_keytab_path_from_base64_keytab_with_decode_exception(self):
+        hook = SparkSubmitHook()
+        invalid_base64 = "invalid_base64"
+
+        with pytest.raises(AirflowException, match="Failed to decode base64 keytab"):
+            hook._create_keytab_path_from_base64_keytab(invalid_base64, None)
+
     @patch("airflow.providers.apache.spark.hooks.spark_submit.uuid.uuid4")
     @patch("pathlib.Path.resolve")
     @patch("airflow.providers.apache.spark.hooks.spark_submit.shutil.move")

From cd6c35898bafaa233cf9dcb20482fc06ec03782c Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Wed, 4 Sep 2024 00:53:59 +0900
Subject: [PATCH 22/25] added keytab write & move exception test case

---
 .../apache/spark/hooks/test_spark_submit.py   | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py
index ce2dcdfe6de42..8e3bdcbbf7ec9 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -1040,6 +1040,55 @@ def test_create_keytab_path_from_base64_keytab_with_decode_exception(self):
         with pytest.raises(AirflowException, match="Failed to decode base64 keytab"):
             hook._create_keytab_path_from_base64_keytab(invalid_base64, None)
 
+    @patch("pathlib.Path.exists")
+    @patch("builtins.open", new_callable=mock_open)
+    def test_create_keytab_path_from_base64_keytab_with_write_exception(
+        self,
+        mock_open,
+        mock_exists,
+    ):
+        # Given
+        hook = SparkSubmitHook()
+
+        keytab_value = b"abcd"
+        base64_keytab = base64.b64encode(keytab_value).decode("UTF-8")
+        _mock_open = mock_open()
+        _mock_open.write.side_effect = Exception("Write failed")
+        mock_exists.return_value = False
+
+        # When
+        with pytest.raises(AirflowException, match="Failed to save keytab"):
+            hook._create_keytab_path_from_base64_keytab(base64_keytab, None)
+
+        # Then
+        assert mock_exists.call_count == 2  # called twice (before write, after write)
+
+    @patch("airflow.providers.apache.spark.hooks.spark_submit.shutil.move")
+    @patch("pathlib.Path.exists")
+    @patch("builtins.open", new_callable=mock_open)
+    def test_create_keytab_path_from_base64_keytab_with_move_exception(
+        self,
+        mock_open,
+        mock_exists,
+        mock_move,
+    ):
+        # Given
+        hook = SparkSubmitHook()
+
+        keytab_value = b"abcd"
+        base64_keytab = base64.b64encode(keytab_value).decode("UTF-8")
+        mock_exists.return_value = False
+        mock_move.side_effect = Exception("Move failed")
+
+        # When
+        with pytest.raises(AirflowException, match="Failed to save keytab"):
+            hook._create_keytab_path_from_base64_keytab(base64_keytab, None)
+
+        # Then
+        mock_open().write.assert_called_once_with(keytab_value)
+        mock_move.assert_called_once()
+        assert mock_exists.call_count == 2  # called twice (before write, after write)
+
     @patch("airflow.providers.apache.spark.hooks.spark_submit.uuid.uuid4")
     @patch("pathlib.Path.resolve")
     @patch("airflow.providers.apache.spark.hooks.spark_submit.shutil.move")

From 55007bc3401eb42bff86dfdd4cb442c0e66c4f4b Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Sat, 7 Sep 2024 08:00:13 +0900
Subject: [PATCH 23/25] removed _get_keytab & _get_principal

---
 .../apache/spark/hooks/spark_submit.py        | 22 +++++--------------
 1 file changed, 6 insertions(+), 16 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index dc2bce3e36ece..d19dfc37e391d 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -356,14 +356,6 @@ def _get_spark_binary_path(self) -> list[str]:
         # Assume that spark-submit is present in the path to the executing user
         return [self._connection["spark_binary"]]
 
-    def _get_principal(self) -> str | None:
-        # for the case where the principal is not set in the broken connection
-        return self._connection["principal"] or self._principal
-
-    def _get_keytab(self) -> str | None:
-        # for the case where the keytab is not set in the broken connection
-        return self._connection["keytab"] or self._keytab
-
     def _mask_cmd(self, connection_cmd: str | list[str]) -> str:
         # Mask any password related fields in application args with key value pair
         # where key contains password (case insensitive), e.g. HivePassword='abc'
@@ -448,12 +440,10 @@ def _build_spark_submit_command(self, application: str) -> list[str]:
             connection_cmd += ["--executor-memory", self._executor_memory]
         if self._driver_memory:
             connection_cmd += ["--driver-memory", self._driver_memory]
-        keytab = self._get_keytab()
-        if keytab:
-            connection_cmd += ["--keytab", keytab]
-        principal = self._get_principal()
-        if principal:
-            connection_cmd += ["--principal", principal]
+        if self._connection["keytab"]:
+            connection_cmd += ["--keytab", self._connection["keytab"]]
+        if self._connection["principal"]:
+            connection_cmd += ["--principal", self._connection["principal"]]
         if self._use_krb5ccache:
             if not os.getenv("KRB5CCNAME"):
                 raise AirflowException(
@@ -788,8 +778,8 @@ def on_kill(self) -> None:
             if self._yarn_application_id:
                 kill_cmd = f"yarn application -kill {self._yarn_application_id}".split()
                 env = {**os.environ, **(self._env or {})}
-                principal = self._get_principal()
-                keytab = self._get_keytab()
+                principal = self._connection["principal"]
+                keytab = self._connection["keytab"]
                 if keytab is not None and principal is not None:
                     # we are ignoring renewal failures from renew_from_kt
                     # here as the failure could just be due to a non-renewable ticket,

From 94aa9266aa36bdc6a1e1ae364968b9ecee831a3d Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Mon, 16 Sep 2024 22:40:01 +0900
Subject: [PATCH 24/25] rollback _get_keytab & _get_principal

---
 .../apache/spark/hooks/spark_submit.py        | 22 ++++++++++++++-----
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index d19dfc37e391d..dc2bce3e36ece 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -356,6 +356,14 @@ def _get_spark_binary_path(self) -> list[str]:
         # Assume that spark-submit is present in the path to the executing user
         return [self._connection["spark_binary"]]
 
+    def _get_principal(self) -> str | None:
+        # for the case where the principal is not set in the broken connection
+        return self._connection["principal"] or self._principal
+
+    def _get_keytab(self) -> str | None:
+        # for the case where the keytab is not set in the broken connection
+        return self._connection["keytab"] or self._keytab
+
     def _mask_cmd(self, connection_cmd: str | list[str]) -> str:
         # Mask any password related fields in application args with key value pair
         # where key contains password (case insensitive), e.g. HivePassword='abc'
@@ -440,10 +448,12 @@ def _build_spark_submit_command(self, application: str) -> list[str]:
             connection_cmd += ["--executor-memory", self._executor_memory]
         if self._driver_memory:
             connection_cmd += ["--driver-memory", self._driver_memory]
-        if self._connection["keytab"]:
-            connection_cmd += ["--keytab", self._connection["keytab"]]
-        if self._connection["principal"]:
-            connection_cmd += ["--principal", self._connection["principal"]]
+        keytab = self._get_keytab()
+        if keytab:
+            connection_cmd += ["--keytab", keytab]
+        principal = self._get_principal()
+        if principal:
+            connection_cmd += ["--principal", principal]
         if self._use_krb5ccache:
             if not os.getenv("KRB5CCNAME"):
                 raise AirflowException(
@@ -778,8 +788,8 @@ def on_kill(self) -> None:
             if self._yarn_application_id:
                 kill_cmd = f"yarn application -kill {self._yarn_application_id}".split()
                 env = {**os.environ, **(self._env or {})}
-                principal = self._connection["principal"]
-                keytab = self._connection["keytab"]
+                principal = self._get_principal()
+                keytab = self._get_keytab()
                 if keytab is not None and principal is not None:
                     # we are ignoring renewal failures from renew_from_kt
                     # here as the failure could just be due to a non-renewable ticket,

From 01b22d41b0ee3a7e36afa4348660988ff515c418 Mon Sep 17 00:00:00 2001
From: seyoon-lim <sylim123@korea.ac.kr>
Date: Mon, 16 Sep 2024 23:15:02 +0900
Subject: [PATCH 25/25] set default values of keytab & principal on conn_data

---
 .../apache/spark/hooks/spark_submit.py        | 46 ++++++++-----------
 1 file changed, 18 insertions(+), 28 deletions(-)

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
index dc2bce3e36ece..37f624e310873 100644
--- a/airflow/providers/apache/spark/hooks/spark_submit.py
+++ b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -258,8 +258,8 @@ def _resolve_connection(self) -> dict[str, Any]:
             "deploy_mode": None,
             "spark_binary": self.spark_binary or DEFAULT_SPARK_BINARY,
             "namespace": None,
-            "principal": None,
-            "keytab": None,
+            "principal": self._principal,
+            "keytab": self._keytab,
         }
 
         try:
@@ -292,14 +292,14 @@ def _resolve_connection(self) -> dict[str, Any]:
                 )
             conn_data["spark_binary"] = self.spark_binary
             conn_data["namespace"] = extra.get("namespace")
-            conn_data["principal"] = self._principal or extra.get("principal")
-            base64_keytab = extra.get("keytab")
-            if self._keytab is not None:
-                conn_data["keytab"] = self._keytab
-            elif base64_keytab is not None:
-                conn_data["keytab"] = self._create_keytab_path_from_base64_keytab(
-                    base64_keytab, conn_data["principal"]
-                )
+            if conn_data["principal"] is None:
+                conn_data["principal"] = extra.get("principal")
+            if conn_data["keytab"] is None:
+                base64_keytab = extra.get("keytab")
+                if base64_keytab is not None:
+                    conn_data["keytab"] = self._create_keytab_path_from_base64_keytab(
+                        base64_keytab, conn_data["principal"]
+                    )
         except AirflowException:
             self.log.info(
                 "Could not load connection string %s, defaulting to %s", self._conn_id, conn_data["master"]
@@ -356,14 +356,6 @@ def _get_spark_binary_path(self) -> list[str]:
         # Assume that spark-submit is present in the path to the executing user
         return [self._connection["spark_binary"]]
 
-    def _get_principal(self) -> str | None:
-        # for the case where the principal is not set in the broken connection
-        return self._connection["principal"] or self._principal
-
-    def _get_keytab(self) -> str | None:
-        # for the case where the keytab is not set in the broken connection
-        return self._connection["keytab"] or self._keytab
-
     def _mask_cmd(self, connection_cmd: str | list[str]) -> str:
         # Mask any password related fields in application args with key value pair
         # where key contains password (case insensitive), e.g. HivePassword='abc'
@@ -448,12 +440,10 @@ def _build_spark_submit_command(self, application: str) -> list[str]:
             connection_cmd += ["--executor-memory", self._executor_memory]
         if self._driver_memory:
             connection_cmd += ["--driver-memory", self._driver_memory]
-        keytab = self._get_keytab()
-        if keytab:
-            connection_cmd += ["--keytab", keytab]
-        principal = self._get_principal()
-        if principal:
-            connection_cmd += ["--principal", principal]
+        if self._connection["keytab"]:
+            connection_cmd += ["--keytab", self._connection["keytab"]]
+        if self._connection["principal"]:
+            connection_cmd += ["--principal", self._connection["principal"]]
         if self._use_krb5ccache:
             if not os.getenv("KRB5CCNAME"):
                 raise AirflowException(
@@ -788,13 +778,13 @@ def on_kill(self) -> None:
             if self._yarn_application_id:
                 kill_cmd = f"yarn application -kill {self._yarn_application_id}".split()
                 env = {**os.environ, **(self._env or {})}
-                principal = self._get_principal()
-                keytab = self._get_keytab()
-                if keytab is not None and principal is not None:
+                if self._connection["keytab"] is not None and self._connection["principal"] is not None:
                     # we are ignoring renewal failures from renew_from_kt
                     # here as the failure could just be due to a non-renewable ticket,
                     # we still attempt to kill the yarn application
-                    renew_from_kt(principal, keytab, exit_on_fail=False)
+                    renew_from_kt(
+                        self._connection["principal"], self._connection["keytab"], exit_on_fail=False
+                    )
                     env = os.environ.copy()
                     ccacche = airflow_conf.get_mandatory_value("kerberos", "ccache")
                     env["KRB5CCNAME"] = ccacche
