diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py index 199d6a6d35dd3..5a55c7d04b70a 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py @@ -805,26 +805,35 @@ def extract_xcom(self, pod: V1Pod) -> str: ) def extract_xcom_json(self, pod: V1Pod) -> str: """Retrieve XCom value and also check if xcom json is valid.""" + command = ( + f"if [ -s {PodDefaults.XCOM_MOUNT_PATH}/return.json ]; " + f"then cat {PodDefaults.XCOM_MOUNT_PATH}/return.json; " + f"else echo {EMPTY_XCOM_RESULT}; fi" + ) with closing( kubernetes_stream( self._client.connect_get_namespaced_pod_exec, pod.metadata.name, pod.metadata.namespace, container=PodDefaults.SIDECAR_CONTAINER_NAME, - command=["/bin/sh"], - stdin=True, + command=[ + "/bin/sh", + "-c", + command, + ], + stdin=False, stdout=True, stderr=True, tty=False, _preload_content=False, ) - ) as resp: - result = self._exec_pod_command( - resp, - f"if [ -s {PodDefaults.XCOM_MOUNT_PATH}/return.json ]; " - f"then cat {PodDefaults.XCOM_MOUNT_PATH}/return.json; " - f"else echo {EMPTY_XCOM_RESULT}; fi", - ) + ) as client: + self.log.info("Running command... %s", command) + client.run_forever() + if client.peek_stderr(): + stderr = client.read_stderr() + self.log.error("stderr from command: %s", stderr) + result = client.read_all() if result and result.rstrip() != EMPTY_XCOM_RESULT: # Note: result string is parsed to check if its valid json. # This function still returns a string which is converted into json in the calling method. diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py index bda48af1eebcc..4091e1e30e28d 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py @@ -539,48 +539,56 @@ def test_container_is_terminated_with_waiting_state(self, container_state, expec assert container_is_terminated(pod_info, "base") == expected_is_terminated @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.kubernetes_stream") - @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager._exec_pod_command") @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.extract_xcom_kill") - def test_extract_xcom_success(self, mock_exec_xcom_kill, mock_exec_pod_command, mock_kubernetes_stream): + def test_extract_xcom_success(self, mock_exec_xcom_kill, mock_kubernetes_stream): """test when valid json is retrieved from xcom sidecar container.""" xcom_json = """{"a": "true"}""" mock_pod = MagicMock() - mock_exec_pod_command.return_value = xcom_json + mock_client = MagicMock() + mock_client.peek_stderr.return_value = "" + mock_client.read_all.return_value = xcom_json + mock_kubernetes_stream.return_value = mock_client ret = self.pod_manager.extract_xcom(pod=mock_pod) assert ret == xcom_json assert mock_exec_xcom_kill.call_count == 1 @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.kubernetes_stream") - @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager._exec_pod_command") @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.extract_xcom_kill") - def test_extract_xcom_failure(self, mock_exec_xcom_kill, mock_exec_pod_command, mock_kubernetes_stream): + def test_extract_xcom_failure(self, mock_exec_xcom_kill, mock_kubernetes_stream): """test when invalid json is retrieved from xcom sidecar container.""" xcom_json = """{"a": "tru""" # codespell:ignore tru mock_pod = MagicMock() - mock_exec_pod_command.return_value = xcom_json + mock_client = MagicMock() + mock_client.peek_stderr.return_value = "" + mock_client.read_all.return_value = xcom_json + mock_kubernetes_stream.return_value = mock_client with pytest.raises(JSONDecodeError): self.pod_manager.extract_xcom(pod=mock_pod) assert mock_exec_xcom_kill.call_count == 1 @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.kubernetes_stream") - @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager._exec_pod_command") @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.extract_xcom_kill") - def test_extract_xcom_empty(self, mock_exec_xcom_kill, mock_exec_pod_command, mock_kubernetes_stream): + def test_extract_xcom_empty(self, mock_exec_xcom_kill, mock_kubernetes_stream): """test when __airflow_xcom_result_empty__ is retrieved from xcom sidecar container.""" mock_pod = MagicMock() xcom_result = "__airflow_xcom_result_empty__" - mock_exec_pod_command.return_value = xcom_result + mock_client = MagicMock() + mock_client.peek_stderr.return_value = "" + mock_client.read_all.return_value = xcom_result + mock_kubernetes_stream.return_value = mock_client ret = self.pod_manager.extract_xcom(pod=mock_pod) assert ret == xcom_result assert mock_exec_xcom_kill.call_count == 1 @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.kubernetes_stream") - @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager._exec_pod_command") @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.extract_xcom_kill") - def test_extract_xcom_none(self, mock_exec_xcom_kill, mock_exec_pod_command, mock_kubernetes_stream): + def test_extract_xcom_none(self, mock_exec_xcom_kill, mock_kubernetes_stream): """test when None is retrieved from xcom sidecar container.""" mock_pod = MagicMock() - mock_exec_pod_command.return_value = None + mock_client = MagicMock() + mock_client.peek_stderr.return_value = "" + mock_client.read_all.return_value = None + mock_kubernetes_stream.return_value = mock_client with pytest.raises(AirflowException): self.pod_manager.extract_xcom(pod=mock_pod) assert mock_exec_xcom_kill.call_count == 1