ferruzzi commented on code in PR #39245:
URL: https://github.com/apache/airflow/pull/39245#discussion_r1585546006
##########
airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -211,3 +262,149 @@ def execute(self, context: Context) -> Any:
)
else:
super().execute(context=context)
+
+
+class BedrockKnowledgeBaseActiveSensor(BedrockAgentBaseSensor):
+ """
+ Poll the Knowledge Base status until it reaches a terminal state; fails if
creation fails.
+
+ .. seealso::
+ For more information on how to use this sensor, take a look at the
guide:
+ :ref:`howto/sensor:BedrockKnowledgeBaseActiveSensor`
+
+ :param knowledge_base_id: The unique identifier of the knowledge base for
which to get information. (templated)
+
+ :param deferrable: If True, the sensor will operate in deferrable more.
This mode requires aiobotocore
+ module to be installed.
+ (default: False, but can be overridden in config file by setting
default_deferrable to True)
+ :param poke_interval: Polling period in seconds to check for the status of
the job. (default: 5)
+ :param max_retries: Number of times before returning the current state
(default: 24)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used.
If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default
boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param botocore_config: Configuration dictionary (key-values) for botocore
client. See:
+
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ """
+
+ INTERMEDIATE_STATES: tuple[str, ...] = ("CREATING", "UPDATING")
+ FAILURE_STATES: tuple[str, ...] = ("DELETING", "FAILED")
+ SUCCESS_STATES: tuple[str, ...] = ("ACTIVE",)
+ FAILURE_MESSAGE = "Bedrock Knowledge Base Active sensor failed."
+
+ template_fields: Sequence[str] = aws_template_fields("knowledge_base_id")
+
+ def __init__(
+ self,
+ *,
+ knowledge_base_id: str,
+ poke_interval: int = 5,
+ max_retries: int = 24,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.poke_interval = poke_interval
+ self.max_retries = max_retries
+ self.knowledge_base_id = knowledge_base_id
+
+ def get_state(self) -> str:
+ return
self.hook.conn.get_knowledge_base(knowledgeBaseId=self.knowledge_base_id)["knowledgeBase"][
+ "status"
+ ]
+
+ def execute(self, context: Context) -> Any:
+ if self.deferrable:
+ self.defer(
+ trigger=BedrockKnowledgeBaseActiveTrigger(
+ knowledge_base_id=self.knowledge_base_id,
+ waiter_delay=int(self.poke_interval),
+ waiter_max_attempts=self.max_retries,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="poke",
Review Comment:
if I modify poke from
```
def poke(self, context: Context, **kwargs) -> bool:
state = self.get_state()
if state in self.FAILURE_STATES:
# TODO: remove this if block when min_airflow_version is set to
higher than 2.7.1
if self.soft_fail:
raise AirflowSkipException(self.FAILURE_MESSAGE)
raise AirflowException(self.FAILURE_MESSAGE)
return state not in self.INTERMEDIATE_STATES
```
to (untested)
```
def poke(self, context: Context, **kwargs) -> bool:
state = kwargs.gget("event").get("state") or self.get_state()
if state in self.FAILURE_STATES or state == "failure":
# TODO: remove this if block when min_airflow_version is set to
higher than 2.7.1
if self.soft_fail:
raise AirflowSkipException(self.FAILURE_MESSAGE)
raise AirflowException(self.FAILURE_MESSAGE)
return state not in self.INTERMEDIATE_STATES
```
Then that would be the same thing without the extra get_state() call
--
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]
