vincbeck commented on code in PR #38850:
URL: https://github.com/apache/airflow/pull/38850#discussion_r1557596522
##########
airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -17,36 +17,89 @@
# under the License.
from __future__ import annotations
+import abc
from typing import TYPE_CHECKING, Any, Sequence
from airflow.configuration import conf
+from airflow.exceptions import AirflowException, AirflowSkipException
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
-from airflow.providers.amazon.aws.triggers.bedrock import
BedrockCustomizeModelCompletedTrigger
+from airflow.providers.amazon.aws.triggers.bedrock import (
+ BedrockCustomizeModelCompletedTrigger,
+ BedrockProvisionModelThroughputCompletedTrigger,
+)
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
if TYPE_CHECKING:
from airflow.utils.context import Context
-from airflow.exceptions import AirflowException, AirflowSkipException
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+class BaseBedrockSensor(AwsBaseSensor[BedrockHook]):
+ """
+ General sensor behavior for Amazon Bedrock.
+
+ Subclasses must implement following methods:
+ - ``get_state()``
-class BedrockCustomizeModelCompletedSensor(AwsBaseSensor[BedrockHook]):
+ Subclasses must set the following fields:
+ - ``INTERMEDIATE_STATES``
+ - ``FAILURE_STATES`
+ - ``SUCCESS_STATES``
+ - ``FAILURE_MESSAGE``
+
+ :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).
+ """
+
+ INTERMEDIATE_STATES: tuple[str, ...] = ()
+ FAILURE_STATES: tuple[str, ...] = ()
+ SUCCESS_STATES: tuple[str, ...] = ()
+ FAILURE_MESSAGE = ""
+
+ aws_hook_class = BedrockHook
+ ui_color = "#66c3ff"
+
+ def __init__(
+ self,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
Review Comment:
nit: missing from docstring
--
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]
