ferruzzi commented on code in PR #39245:
URL: https://github.com/apache/airflow/pull/39245#discussion_r1585748958
##########
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:
Yeah, I initially added it as an explicit parameter but it raised issues
with the static checks for not matching the signature of super(). Adding
`**kwargs` passed the tests and `event` will just get swallowed up in there.
Since I'm not actually using it, I thought it was acceptable and in the end I
decided it actually looked less confusing this way (as opposed to seeing
'event" there and wondering what you are supposed to give it if you aren't
familiar with deferrables)
Are you alright with it how it is? Maybe there's another way around that?
I can look if you care. If you are alright with it how it is, can you resolve
this convo?
--
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]
