syedahsn commented on code in PR #39245:
URL: https://github.com/apache/airflow/pull/39245#discussion_r1585669891
##########
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:
That makes sense to me. I think the extra call is worth the clarity. One
thing I noticed is that `poke` doesn't accept `event` as a parameter. I'm not
sure how important it is the the functionality (I'm sure you've tested it), but
the docs seem to strongly recommend it:
>When your operator resumes, Airflow adds a context object and an event
object to the kwargs passed to the method_name method. This event object
contains the payload from the trigger event that resumed your operator.
Depending on the trigger, this can be useful to your operator, like it’s a
status code or URL to fetch results. Or, it might be unimportant information,
like a datetime. Your method_name method, however, _must_ accept context and
event as a keyword argument.
I don't feel too strongly about it either way.
LGTM
--
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]
