diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index 709583d78e2ed..5b9e57ec01743 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -80,7 +80,6 @@ PodNotFoundException, PodOperatorHookProtocol, PodPhase, - check_exception_is_kubernetes_api_unauthorized, container_is_succeeded, get_container_termination_message, ) @@ -113,6 +112,10 @@ class PodReattachFailure(AirflowException): """When we expect to be able to find a pod but cannot.""" +class PodCredentialsExpiredFailure(AirflowException): + """When pod fails to refresh credentials.""" + + class KubernetesPodOperator(BaseOperator): """ Execute a task in a Kubernetes Pod. @@ -652,9 +655,8 @@ def execute_sync(self, context: Context): return result @tenacity.retry( - stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_exponential(max=15), - retry=tenacity.retry_if_exception(lambda exc: check_exception_is_kubernetes_api_unauthorized(exc)), + retry=tenacity.retry_if_exception_type(PodCredentialsExpiredFailure), reraise=True, ) def await_pod_completion(self, pod: k8s.V1Pod): @@ -675,6 +677,10 @@ def await_pod_completion(self, pod: k8s.V1Pod): "Failed to check container status due to permission error. Refreshing credentials and retrying." ) self._refresh_cached_properties() + self.pod_manager.read_pod( + pod=pod + ) # attempt using refreshed credentials, raises if still invalid + raise PodCredentialsExpiredFailure("Kubernetes credentials expired, retrying after refresh.") raise exc def _refresh_cached_properties(self): diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index 1ca0d9851b18c..be8279bcabd9a 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -1629,8 +1629,9 @@ def test_execute_async_callbacks(self): @pytest.mark.parametrize("get_logs", [True, False]) @patch(f"{POD_MANAGER_CLASS}.fetch_requested_container_logs") @patch(f"{POD_MANAGER_CLASS}.await_container_completion") + @patch(f"{POD_MANAGER_CLASS}.read_pod") def test_await_container_completion_refreshes_properties_on_exception( - self, mock_await_container_completion, fetch_requested_container_logs, get_logs + self, mock_read_pod, mock_await_container_completion, fetch_requested_container_logs, get_logs ): k = KubernetesPodOperator(task_id="task", get_logs=get_logs) pod = self.run_pod(k) @@ -1655,6 +1656,28 @@ def test_await_container_completion_refreshes_properties_on_exception( mock_await_container_completion.assert_has_calls( [mock.call(pod=pod, container_name=k.base_container_name)] * 3 ) + mock_read_pod.assert_called() + assert client != k.client + assert hook != k.hook + assert pod_manager != k.pod_manager + + @patch(f"{POD_MANAGER_CLASS}.await_container_completion") + @patch(f"{POD_MANAGER_CLASS}.read_pod") + def test_await_container_completion_raises_unauthorized_if_credentials_still_invalid_after_refresh( + self, mock_read_pod, mock_await_container_completion + ): + k = KubernetesPodOperator(task_id="task", get_logs=False) + pod = self.run_pod(k) + client, hook, pod_manager = k.client, k.hook, k.pod_manager + + mock_await_container_completion.side_effect = [ApiException(status=401)] + mock_read_pod.side_effect = [ApiException(status=401)] + + with pytest.raises(ApiException): + k.await_pod_completion(pod) + + mock_read_pod.assert_called() + # assert cache was refreshed assert client != k.client assert hook != k.hook assert pod_manager != k.pod_manager @@ -1663,7 +1686,7 @@ def test_await_container_completion_refreshes_properties_on_exception( "side_effect, exception_type, expect_exc", [ ([ApiException(401), mock.DEFAULT], ApiException, True), # works after one 401 - ([ApiException(401)] * 10, ApiException, False), # exc after 3 retries on 401 + ([ApiException(401)] * 3 + [mock.DEFAULT], ApiException, True), # works after 3 retries ([ApiException(402)], ApiException, False), # exc on non-401 ([ApiException(500)], ApiException, False), # exc on non-401 ([Exception], Exception, False), # exc on different exception @@ -1684,7 +1707,7 @@ def test_await_container_completion_retries_on_specific_exception( else: with pytest.raises(exception_type): k.await_pod_completion(pod) - expected_call_count = min(len(side_effect), 3) # retry max 3 times + expected_call_count = len(side_effect) mock_await_container_completion.assert_has_calls( [mock.call(pod=pod, container_name=k.base_container_name)] * expected_call_count )