Taragolis commented on code in PR #39245:
URL: https://github.com/apache/airflow/pull/39245#discussion_r1585162253
##########
airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -68,7 +70,56 @@ def __init__(
super().__init__(**kwargs)
self.deferrable = deferrable
- def poke(self, context: Context) -> bool:
+ 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
+
+ @abc.abstractmethod
+ def get_state(self) -> str:
+ """Implement in subclasses."""
+
+
Review Comment:
If you would like to make generic from the generic you could do something
like that
```python
_GenericBedrockHook = TypeVar("_GenericBedrockHook", BedrockAgentHook,
BedrockHook)
class BedrockBaseSensor(AwsBaseSensor[_GenericBedrockHook]):
"""
General sensor behavior for Amazon Bedrock.
Subclasses must implement following methods:
- ``get_state()``
Subclasses must set the following fields:
- ``INTERMEDIATE_STATES``
- ``FAILURE_STATES``
- ``SUCCESS_STATES``
- ``FAILURE_MESSAGE``
:param deferrable: If True, the sensor will operate in deferrable mode.
This mode requires aiobotocore
module to be installed.
(default: False, but can be overridden in config file by setting
default_deferrable to True)
"""
INTERMEDIATE_STATES: tuple[str, ...]
FAILURE_STATES: tuple[str, ...]
SUCCESS_STATES: tuple[str, ...]
FAILURE_MESSAGE: str
ui_color = "#66c3ff"
aws_hook_class: type[_GenericBedrockHook]
...
class BedrockCustomizeModelCompletedSensor(BedrockBaseSensor[BedrockHook]):
aws_hook_class = BedrockHook
...
class BedrockKnowledgeBaseActiveSensor(BedrockBaseSensor[BedrockAgentHook]):
aws_hook_class = BedrockAgentHook
...
```
--
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]
