[airflow] branch main updated: Refactor: Extract and reuse get_kerberos_principle func from get_kerberos_principle (#34936)

[eladkal]

This is an automated email from the ASF dual-hosted git repository.

eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 927b0bbe94 Refactor: Extract and reuse get_kerberos_principle func 
from get_kerberos_principle (#34936)
927b0bbe94 is described below

commit 927b0bbe94ed97c395944da23ab042b02e701af1
Author: zeotuan <48720253+zeot...@users.noreply.github.com>
AuthorDate: Thu Oct 19 07:53:51 2023 +1100

    Refactor: Extract and reuse get_kerberos_principle func from 
get_kerberos_principle (#34936)
    
    * Refactor: Extract get_kerberos_principle func
    
    * Refactor: Extract get_kerberos_principle func
    
    * Add tests get_kerberos_principle
    
    * Add doc String
    
    ---------
    
    Co-authored-by: Elad Kalif <45845474+elad...@users.noreply.github.com>
---
 airflow/security/kerberos.py    | 10 ++++++----
 tests/security/test_kerberos.py | 13 ++++++++++++-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/airflow/security/kerberos.py b/airflow/security/kerberos.py
index 32a62f13de..69afc5d793 100644
--- a/airflow/security/kerberos.py
+++ b/airflow/security/kerberos.py
@@ -47,6 +47,11 @@ NEED_KRB181_WORKAROUND: bool | None = None
 log = logging.getLogger(__name__)
 
 
+def get_kerberos_principle(principal: str | None) -> str:
+    """Retrieve Kerberos principal. Fallback to principal from Airflow 
configuration if not provided."""
+    return principal or conf.get_mandatory_value("kerberos", 
"principal").replace("_HOST", get_hostname())
+
+
 def renew_from_kt(principal: str | None, keytab: str, exit_on_fail: bool = 
True):
     """
     Renew kerberos token from keytab.
@@ -59,10 +64,7 @@ def renew_from_kt(principal: str | None, keytab: str, 
exit_on_fail: bool = True)
     # minutes to give ourselves a large renewal buffer.
     renewal_lifetime = f"{conf.getint('kerberos', 'reinit_frequency')}m"
 
-    cmd_principal = principal or conf.get_mandatory_value("kerberos", 
"principal").replace(
-        "_HOST", get_hostname()
-    )
-
+    cmd_principal = get_kerberos_principle(principal)
     if conf.getboolean("kerberos", "forwardable"):
         forwardable = "-f"
     else:
diff --git a/tests/security/test_kerberos.py b/tests/security/test_kerberos.py
index 12dfe93dd6..50da424c3d 100644
--- a/tests/security/test_kerberos.py
+++ b/tests/security/test_kerberos.py
@@ -24,7 +24,7 @@ from unittest import mock
 import pytest
 
 from airflow.security import kerberos
-from airflow.security.kerberos import renew_from_kt
+from airflow.security.kerberos import get_kerberos_principle, renew_from_kt
 from tests.test_utils.config import conf_vars
 
 
@@ -281,3 +281,14 @@ class TestKerberos:
             mock.call("test-principal", "/tmp/keytab"),
             mock.call("test-principal", "/tmp/keytab"),
         ]
+
+    def test_get_kerberos_principle(self):
+        expected_principal = "test-principal"
+        principal = get_kerberos_principle(expected_principal)
+        assert principal == expected_principal
+
+    @mock.patch("airflow.security.kerberos.get_hostname", 
return_value="REPLACEMENT_HOST")
+    @mock.patch("airflow.security.kerberos.conf.get_mandatory_value", 
return_value="test-principal/_HOST")
+    def test_get_kerberos_principle_resolve_null_principal(self, 
get_madantory_value_mock, get_hostname_mock):
+        principal = get_kerberos_principle(principal=None)
+        assert principal == "test-principal/REPLACEMENT_HOST"