o-nikolas commented on code in PR #55445:
URL: https://github.com/apache/airflow/pull/55445#discussion_r2337443949
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -654,9 +656,37 @@ def execute_complete(self, context: Context, event:
dict[str, Any] | None = None
return validated_event["ingestion_job_id"]
def execute(self, context: Context) -> str:
- ingestion_job_id = self.hook.conn.start_ingestion_job(
- knowledgeBaseId=self.knowledge_base_id,
dataSourceId=self.data_source_id
- )["ingestionJob"]["ingestionJobId"]
+ def start_ingestion_job():
+ try:
+ ingestion_job_id = self.hook.conn.start_ingestion_job(
+ knowledgeBaseId=self.knowledge_base_id,
dataSourceId=self.data_source_id
+ )["ingestionJob"]["ingestionJobId"]
+
+ return ingestion_job_id
+ except ClientError as error:
+ error_message = error.response["Error"]["Message"].lower()
+ is_known_retryable_message = (
Review Comment:
Okay, I have this tested and working and I don't want to burn too much more
time on this refactoring it right now. I think I'll merge it so that we get the
benefit and we can circle back and change all the operators to use tenacity
later if we find the time!
Thanks for the review :smiley:
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
