This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 3260046a46e Include Bedrock ingestion job failure reasons in the error
message (#71124)
3260046a46e is described below
commit 3260046a46ed100591284d4fdecc53a93e38f31e
Author: D. Ferruzzi <[email protected]>
AuthorDate: Mon Aug 10 08:03:19 2026 -0700
Include Bedrock ingestion job failure reasons in the error message (#71124)
---
.../providers/amazon/aws/sensors/bedrock.py | 27 +++++++++++++++++++---
.../tests/unit/amazon/aws/sensors/test_bedrock.py | 18 ++++++++++++---
2 files changed, 39 insertions(+), 6 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py
b/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py
index 7ccac1f87eb..5d00918ec46 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py
@@ -79,14 +79,27 @@ class BedrockBaseSensor(AwsBaseSensor[_GenericBedrockHook]):
def poke(self, context: Context, **kwargs) -> bool:
state = self.get_state()
if state in self.FAILURE_STATES:
- raise AirflowException(self.FAILURE_MESSAGE)
+ raise
AirflowException(f"{self.FAILURE_MESSAGE}{self._failure_reason_suffix()}")
return state not in self.INTERMEDIATE_STATES
+ def _failure_reason_suffix(self) -> str:
+ """Assemble failure reasons if any are provided."""
+ return f" Failure reasons: {reason}" if (reason :=
self.get_failure_reason()) else ""
+
@abc.abstractmethod
def get_state(self) -> str:
"""Implement in subclasses."""
+ def get_failure_reason(self) -> str:
+ """
+ Return the service-reported reason(s) for a failed state, or an empty
string if none.
+
+ Override in subclasses whose describe API exposes failure detail.
Return the reasons
+ only; phrasing and separators are handled by the caller.
+ """
+ return ""
+
class BedrockCustomizeModelCompletedSensor(BedrockBaseSensor[BedrockHook]):
"""
@@ -348,12 +361,20 @@ class
BedrockIngestionJobSensor(BedrockBaseSensor[BedrockAgentHook]):
self.data_source_id = data_source_id
self.ingestion_job_id = ingestion_job_id
- def get_state(self) -> str:
+ def _get_ingestion_job(self) -> dict[str, Any]:
return self.hook.conn.get_ingestion_job(
knowledgeBaseId=self.knowledge_base_id,
ingestionJobId=self.ingestion_job_id,
dataSourceId=self.data_source_id,
- )["ingestionJob"]["status"]
+ )["ingestionJob"]
+
+ def get_state(self) -> str:
+ return self._get_ingestion_job()["status"]
+
+ def get_failure_reason(self) -> str:
+ if reasons := self._get_ingestion_job().get("failureReasons"):
+ return "; ".join(reasons)
+ return ""
def execute(self, context: Context) -> Any:
if self.deferrable:
diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py
b/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py
index 559926579a4..e3e3792b290 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py
@@ -239,12 +239,24 @@ class TestBedrockIngestionJobSensor:
assert self.sensor.poke({}) is False
@pytest.mark.parametrize("state", SENSOR.FAILURE_STATES)
+ @pytest.mark.parametrize(
+ ("job_detail", "expected_suffix"),
+ [
+ pytest.param({}, "", id="no reasons"),
+ pytest.param(
+ {"failureReasons": ["User is not authorized", "index not
found"]},
+ " Failure reasons: User is not authorized; index not found",
+ id="with reasons",
+ ),
+ ],
+ )
@mock.patch.object(BedrockAgentHook, "conn")
- def test_poke_failure_states(self, mock_conn, state):
- mock_conn.get_ingestion_job.return_value = {"ingestionJob": {"status":
state}}
+ def test_poke_failure_states(self, mock_conn, state, job_detail,
expected_suffix):
+ mock_conn.get_ingestion_job.return_value = {"ingestionJob": {"status":
state, **job_detail}}
sensor = self.SENSOR(**self.default_op_kwargs, aws_conn_id=None)
- with pytest.raises(AirflowException, match=sensor.FAILURE_MESSAGE):
+ with pytest.raises(AirflowException, match=sensor.FAILURE_MESSAGE) as
exception_info:
sensor.poke({})
+ assert str(exception_info.value) ==
f"{sensor.FAILURE_MESSAGE}{expected_suffix}"
class TestBedrockBatchInferenceSensor:
