ferruzzi commented on code in PR #39245:
URL: https://github.com/apache/airflow/pull/39245#discussion_r1585541983
##########
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:
I think this makes sense. I could write another method that will check the
event state, but poke is already doing that. It results in one extra API call
(the extra poke) but makes the code much nicer. If you look at the
execute_complete methods int he other sensors, all they do is `print state and
if state != "success" then raise an exception`, that's just a copy/paste subset
of the existing poke method.
--
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]
