Move exception check to utils

apache · May 1, 2024 · 5567b00 · 5567b00
1 parent 3316ff4
commit 5567b00
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -78,6 +78,7 @@
     PodNotFoundException,
     PodOperatorHookProtocol,
     PodPhase,
+    check_exception_is_kubernetes_api_unauthorized,
     container_is_succeeded,
     get_container_termination_message,
 )
@@ -99,12 +100,6 @@
 KUBE_CONFIG_ENV_VAR = "KUBECONFIG"
 
 
-def check_exception_is_kubernetes_api_unauthorized(exc: BaseException):
-    return (
-        isinstance(exc, kubernetes.client.exceptions.ApiException) and exc.status and str(exc.status) == "401"
-    )
-
-
 class PodEventType(Enum):
     """Type of Events emitted by kubernetes pod."""
 

diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -190,6 +190,10 @@ def get_container_termination_message(pod: V1Pod, container_name: str):
         return container_status.state.terminated.message if container_status else None
 
 
+def check_exception_is_kubernetes_api_unauthorized(exc: BaseException):
+    return isinstance(exc, ApiException) and exc.status and str(exc.status) == "401"
+
+
 class PodLaunchTimeoutException(AirflowException):
     """When pod does not leave the ``Pending`` phase within specified timeout."""