apache-airflow-providers-amazon 8.21.0__py3-none-any.whl → 8.22.0__py3-none-any.whl

airflow/providers/amazon/aws/sensors/bedrock.py

@@ -18,14 +18,16 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence, TypeVar
 
 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.hooks.bedrock import BedrockAgentHook, BedrockHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.bedrock import (
     BedrockCustomizeModelCompletedTrigger,
+    BedrockIngestionJobTrigger,
+    BedrockKnowledgeBaseActiveTrigger,
     BedrockProvisionModelThroughputCompletedTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
@@ -34,7 +36,10 @@ if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class BedrockBaseSensor(AwsBaseSensor[BedrockHook]):
+_GenericBedrockHook = TypeVar("_GenericBedrockHook", BedrockAgentHook, BedrockHook)
+
+
+class BedrockBaseSensor(AwsBaseSensor[_GenericBedrockHook]):
     """
     General sensor behavior for Amazon Bedrock.
 
@@ -57,7 +62,7 @@ class BedrockBaseSensor(AwsBaseSensor[BedrockHook]):
     SUCCESS_STATES: tuple[str, ...] = ()
     FAILURE_MESSAGE = ""
 
-    aws_hook_class = BedrockHook
+    aws_hook_class: type[_GenericBedrockHook]
     ui_color = "#66c3ff"
 
     def __init__(
@@ -68,7 +73,7 @@ class BedrockBaseSensor(AwsBaseSensor[BedrockHook]):
         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
@@ -83,7 +88,7 @@ class BedrockBaseSensor(AwsBaseSensor[BedrockHook]):
         """Implement in subclasses."""
 
 
-class BedrockCustomizeModelCompletedSensor(BedrockBaseSensor):
+class BedrockCustomizeModelCompletedSensor(BedrockBaseSensor[BedrockHook]):
     """
     Poll the state of the model customization job until it reaches a terminal state; fails if the job fails.
 
@@ -115,6 +120,8 @@ class BedrockCustomizeModelCompletedSensor(BedrockBaseSensor):
     SUCCESS_STATES: tuple[str, ...] = ("Completed",)
     FAILURE_MESSAGE = "Bedrock model customization job sensor failed."
 
+    aws_hook_class = BedrockHook
+
     template_fields: Sequence[str] = aws_template_fields("job_name")
 
     def __init__(
@@ -148,7 +155,7 @@ class BedrockCustomizeModelCompletedSensor(BedrockBaseSensor):
         return self.hook.conn.get_model_customization_job(jobIdentifier=self.job_name)["status"]
 
 
-class BedrockProvisionModelThroughputCompletedSensor(BedrockBaseSensor):
+class BedrockProvisionModelThroughputCompletedSensor(BedrockBaseSensor[BedrockHook]):
     """
     Poll the provisioned model throughput job until it reaches a terminal state; fails if the job fails.
 
@@ -180,6 +187,8 @@ class BedrockProvisionModelThroughputCompletedSensor(BedrockBaseSensor):
     SUCCESS_STATES: tuple[str, ...] = ("InService",)
     FAILURE_MESSAGE = "Bedrock provision model throughput sensor failed."
 
+    aws_hook_class = BedrockHook
+
     template_fields: Sequence[str] = aws_template_fields("model_id")
 
     def __init__(
@@ -211,3 +220,153 @@ class BedrockProvisionModelThroughputCompletedSensor(BedrockBaseSensor):
             )
         else:
             super().execute(context=context)
+
+
+class BedrockKnowledgeBaseActiveSensor(BedrockBaseSensor[BedrockAgentHook]):
+    """
+    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."
+
+    aws_hook_class = BedrockAgentHook
+
+    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",
+            )
+        else:
+            super().execute(context=context)
+
+
+class BedrockIngestionJobSensor(BedrockBaseSensor[BedrockAgentHook]):
+    """
+    Poll the ingestion job 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:BedrockIngestionJobSensor`
+
+    :param knowledge_base_id: The unique identifier of the knowledge base for which to get information. (templated)
+    :param data_source_id: The unique identifier of the data source in the ingestion job. (templated)
+    :param ingestion_job_id: The unique identifier of the ingestion job. (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: 60)
+    :param max_retries: Number of times before returning the current state (default: 10)
+    :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, ...] = ("STARTING", "IN_PROGRESS")
+    FAILURE_STATES: tuple[str, ...] = ("FAILED",)
+    SUCCESS_STATES: tuple[str, ...] = ("COMPLETE",)
+    FAILURE_MESSAGE = "Bedrock ingestion job sensor failed."
+
+    aws_hook_class = BedrockAgentHook
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "knowledge_base_id", "data_source_id", "ingestion_job_id"
+    )
+
+    def __init__(
+        self,
+        *,
+        knowledge_base_id: str,
+        data_source_id: str,
+        ingestion_job_id: str,
+        poke_interval: int = 60,
+        max_retries: int = 10,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
+        self.knowledge_base_id = knowledge_base_id
+        self.data_source_id = data_source_id
+        self.ingestion_job_id = ingestion_job_id
+
+    def get_state(self) -> str:
+        return self.hook.conn.get_ingestion_job(
+            knowledgeBaseId=self.knowledge_base_id,
+            ingestionJobId=self.ingestion_job_id,
+            dataSourceId=self.data_source_id,
+        )["ingestionJob"]["status"]
+
+    def execute(self, context: Context) -> Any:
+        if self.deferrable:
+            self.defer(
+                trigger=BedrockIngestionJobTrigger(
+                    knowledge_base_id=self.knowledge_base_id,
+                    ingestion_job_id=self.ingestion_job_id,
+                    data_source_id=self.data_source_id,
+                    waiter_delay=int(self.poke_interval),
+                    waiter_max_attempts=self.max_retries,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="poke",
+            )
+        else:
+            super().execute(context=context)

airflow/providers/amazon/aws/sensors/opensearch_serverless.py

@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException, AirflowSkipException
+from airflow.providers.amazon.aws.hooks.opensearch_serverless import OpenSearchServerlessHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.opensearch_serverless import (
+    OpenSearchServerlessCollectionActiveTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.utils.helpers import exactly_one
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class OpenSearchServerlessCollectionActiveSensor(AwsBaseSensor[OpenSearchServerlessHook]):
+    """
+    Poll the state of the Collection until it reaches a terminal state; fails if the query fails.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:OpenSearchServerlessCollectionAvailableSensor`
+
+    :param collection_id: A collection ID. You can't provide a name and an ID in the same request.
+    :param collection_name: A collection name. You can't provide a name and an ID in the same request.
+
+    :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: 10)
+    :param max_retries: Number of times before returning the current state (default: 60)
+    :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 = ("CREATING",)
+    FAILURE_STATES = (
+        "DELETING",
+        "FAILED",
+    )
+    SUCCESS_STATES = ("ACTIVE",)
+    FAILURE_MESSAGE = "OpenSearch Serverless Collection sensor failed"
+
+    aws_hook_class = OpenSearchServerlessHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "collection_id",
+        "collection_name",
+    )
+    ui_color = "#66c3ff"
+
+    def __init__(
+        self,
+        *,
+        collection_id: str | None = None,
+        collection_name: str | None = None,
+        poke_interval: int = 10,
+        max_retries: int = 60,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        if not exactly_one(collection_id is None, collection_name is None):
+            raise AttributeError("Either collection_ids or collection_names must be provided, not both.")
+        self.collection_id = collection_id
+        self.collection_name = collection_name
+
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
+        self.deferrable = deferrable
+
+    def poke(self, context: Context, **kwargs) -> bool:
+        call_args = (
+            {"ids": [str(self.collection_id)]}
+            if self.collection_id
+            else {"names": [str(self.collection_name)]}
+        )
+        state = self.hook.conn.batch_get_collection(**call_args)["collectionDetails"][0]["status"]
+
+        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)
+
+        if state in self.INTERMEDIATE_STATES:
+            return False
+        return True
+
+    def execute(self, context: Context) -> Any:
+        if self.deferrable:
+            self.defer(
+                trigger=OpenSearchServerlessCollectionActiveTrigger(
+                    collection_id=self.collection_id,
+                    collection_name=self.collection_name,
+                    waiter_delay=int(self.poke_interval),
+                    waiter_max_attempts=self.max_retries,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="poke",
+            )
+        else:
+            super().execute(context=context)

airflow/providers/amazon/aws/triggers/bedrock.py

@@ -18,7 +18,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 
 if TYPE_CHECKING:
@@ -61,11 +61,49 @@ class BedrockCustomizeModelCompletedTrigger(AwsBaseWaiterTrigger):
         return BedrockHook(aws_conn_id=self.aws_conn_id)
 
 
+class BedrockKnowledgeBaseActiveTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when a Bedrock Knowledge Base reaches the ACTIVE state.
+
+    :param knowledge_base_id: The unique identifier of the knowledge base for which to get information.
+
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 5)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 24)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *,
+        knowledge_base_id: str,
+        waiter_delay: int = 5,
+        waiter_max_attempts: int = 24,
+        aws_conn_id: str | None = None,
+    ) -> None:
+        super().__init__(
+            serialized_fields={"knowledge_base_id": knowledge_base_id},
+            waiter_name="knowledge_base_active",
+            waiter_args={"knowledgeBaseId": knowledge_base_id},
+            failure_message="Bedrock Knowledge Base creation failed.",
+            status_message="Status of Bedrock Knowledge Base job is",
+            status_queries=["status"],
+            return_key="knowledge_base_id",
+            return_value=knowledge_base_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return BedrockAgentHook(aws_conn_id=self.aws_conn_id)
+
+
 class BedrockProvisionModelThroughputCompletedTrigger(AwsBaseWaiterTrigger):
     """
     Trigger when a provisioned throughput job is complete.
 
     :param provisioned_model_id: The ARN or name of the provisioned throughput.
+
     :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
     :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
     :param aws_conn_id: The Airflow connection used for AWS credentials.
@@ -95,3 +133,52 @@ class BedrockProvisionModelThroughputCompletedTrigger(AwsBaseWaiterTrigger):
 
     def hook(self) -> AwsGenericHook:
         return BedrockHook(aws_conn_id=self.aws_conn_id)
+
+
+class BedrockIngestionJobTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when a Bedrock ingestion job reaches the COMPLETE state.
+
+    :param knowledge_base_id: The unique identifier of the knowledge base for which to get information.
+    :param data_source_id: The unique identifier of the data source in the ingestion job.
+    :param ingestion_job_id: The unique identifier of the ingestion job.
+
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 10)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *,
+        knowledge_base_id: str,
+        data_source_id: str,
+        ingestion_job_id: str,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 10,
+        aws_conn_id: str | None = None,
+    ) -> None:
+        super().__init__(
+            serialized_fields={
+                "knowledge_base_id": knowledge_base_id,
+                "data_source_id": data_source_id,
+                "ingestion_job_id": ingestion_job_id,
+            },
+            waiter_name="ingestion_job_complete",
+            waiter_args={
+                "knowledgeBaseId": knowledge_base_id,
+                "dataSourceId": data_source_id,
+                "ingestionJobId": ingestion_job_id,
+            },
+            failure_message="Bedrock ingestion job creation failed.",
+            status_message="Status of Bedrock ingestion job is",
+            status_queries=["status"],
+            return_key="ingestion_job_id",
+            return_value=ingestion_job_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return BedrockAgentHook(aws_conn_id=self.aws_conn_id)

airflow/providers/amazon/aws/triggers/opensearch_serverless.py

@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.amazon.aws.hooks.opensearch_serverless import OpenSearchServerlessHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+from airflow.utils.helpers import exactly_one
+
+if TYPE_CHECKING:
+    from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+
+class OpenSearchServerlessCollectionActiveTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when an Amazon OpenSearch Serverless Collection reaches the ACTIVE state.
+
+    :param collection_id: A collection ID. You can't provide a name and an ID in the same request.
+    :param collection_name: A collection name. You can't provide a name and an ID in the same request.
+
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 20)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *,
+        collection_id: str | None = None,
+        collection_name: str | None = None,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 20,
+        aws_conn_id: str | None = None,
+    ) -> None:
+        if not exactly_one(collection_id is None, collection_name is None):
+            raise AttributeError("Either collection_ids or collection_names must be provided, not both.")
+
+        super().__init__(
+            serialized_fields={"collection_id": collection_id, "collection_name": collection_name},
+            waiter_name="collection_available",
+            waiter_args={"ids": [collection_id]} if collection_id else {"names": [collection_name]},
+            failure_message="OpenSearch Serverless Collection creation failed.",
+            status_message="Status of OpenSearch Serverless Collection is",
+            status_queries=["status"],
+            return_key="collection_id" if collection_id else "collection_name",
+            return_value=collection_id if collection_id else collection_name,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return OpenSearchServerlessHook(aws_conn_id=self.aws_conn_id)

airflow/providers/amazon/aws/waiters/bedrock-agent.json

@@ -0,0 +1,73 @@
+{
+  "version": 2,
+  "waiters": {
+    "knowledge_base_active": {
+      "delay": 5,
+      "maxAttempts": 24,
+      "operation": "getKnowledgeBase",
+      "acceptors": [
+        {
+          "matcher": "path",
+          "argument": "knowledgeBase.status",
+          "expected": "ACTIVE",
+          "state": "success"
+        },
+        {
+          "matcher": "path",
+          "argument": "knowledgeBase.status",
+          "expected": "CREATING",
+          "state": "retry"
+        },
+        {
+          "matcher": "path",
+          "argument": "knowledgeBase.status",
+          "expected": "DELETING",
+          "state": "failure"
+        },
+        {
+          "matcher": "path",
+          "argument": "knowledgeBase.status",
+          "expected": "UPDATING",
+          "state": "retry"
+        },
+        {
+          "matcher": "path",
+          "argument": "knowledgeBase.status",
+          "expected": "FAILED",
+          "state": "failure"
+        }
+      ]
+    },
+    "ingestion_job_complete": {
+      "delay": 60,
+      "maxAttempts": 10,
+      "operation": "getIngestionJob",
+      "acceptors": [
+        {
+          "matcher": "path",
+          "argument": "ingestionJob.status",
+          "expected": "COMPLETE",
+          "state": "success"
+        },
+        {
+          "matcher": "path",
+          "argument": "ingestionJob.status",
+          "expected": "STARTING",
+          "state": "retry"
+        },
+        {
+          "matcher": "path",
+          "argument": "ingestionJob.status",
+          "expected": "IN_PROGRESS",
+          "state": "retry"
+        },
+        {
+          "matcher": "path",
+          "argument": "ingestionJob.status",
+          "expected": "FAILED",
+          "state": "failure"
+        }
+      ]
+    }
+  }
+}

airflow/providers/amazon/aws/waiters/opensearchserverless.json

@@ -0,0 +1,36 @@
+{
+    "version": 2,
+    "waiters": {
+        "collection_available": {
+            "operation": "BatchGetCollection",
+            "delay": 10,
+            "maxAttempts": 120,
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "collectionDetails[0].status",
+                    "expected": "ACTIVE",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "collectionDetails[0].status",
+                    "expected": "DELETING",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "collectionDetails[0].status",
+                    "expected": "CREATING",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "collectionDetails[0].status",
+                    "expected": "FAILED",
+                    "state": "failure"
+                }
+            ]
+        }
+    }
+}