AOSS collection_available waiter/sensor/trigger with tests

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

@@ -0,0 +1,110 @@
+# 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.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.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: 60)
+    :param max_retries: Number of times before returning the current state (default: 20)
+    :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"
+    INVALID_ARGS_MESSAGE = "Either collection_ids or collection_names must be provided, not both."
+
+    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 = 120,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        if not exactly_one(collection_id is None, collection_name is None):
+            raise AttributeError(self.INVALID_ARGS_MESSAGE)
+        self.collection_id = collection_id
+        self.collection_name = collection_name
+
+        self.call_args = (
+            {"ids": [str(self.collection_id)]}
+            if self.collection_id
+            else {"names": [str(self.collection_name)]}
+        )
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
+
+    def poke(self, context: Context) -> bool:
+        collections = self.hook.conn.batch_get_collection(**self.call_args)
+        state = collections["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

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

@@ -0,0 +1,73 @@
+# 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.sensors.opensearch_serverless import (
+    OpenSearchServerlessCollectionActiveSensor,
+)
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+from airflow.utils.helpers import exactly_one, prune_dict
+
+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(OpenSearchServerlessCollectionActiveSensor.INVALID_ARGS_MESSAGE)
+
+        call_args = prune_dict({"ids": [collection_id], "names": [collection_name]})
+
+        super().__init__(
+            serialized_fields={"collection_id": collection_id, "collection_name": collection_name},
+            waiter_name="collection_available",
+            waiter_args=call_args,
+            failure_message="OpenSearch Serverless Collection creation failed.",
+            status_message="Status of OpenSearch Serverless Collection is",
+            status_queries=["status"],
+            return_key="collection_id" if "ids" in call_args.keys() else "collection_name",
+            return_value=collection_id if "ids" in call_args.keys() 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/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"
+                }
+            ]
+        }
+    }
+}

airflow/providers/amazon/provider.yaml

@@ -501,6 +501,9 @@ sensors:
   - integration-name: AWS Lambda
     python-modules:
       - airflow.providers.amazon.aws.sensors.lambda_function
+  - integration-name: Amazon OpenSearch Serverless
+    python-modules:
+      - airflow.providers.amazon.aws.sensors.opensearch_serverless
   - integration-name: Amazon RDS
     python-modules:
       - airflow.providers.amazon.aws.sensors.rds
@@ -674,6 +677,9 @@ triggers:
   - integration-name: AWS Lambda
     python-modules:
       - airflow.providers.amazon.aws.triggers.lambda_function
+  - integration-name: Amazon OpenSearch Serverless
+    python-modules:
+      - airflow.providers.amazon.aws.triggers.opensearch_serverless
   - integration-name: Amazon Redshift
     python-modules:
       - airflow.providers.amazon.aws.triggers.redshift_cluster

tests/providers/amazon/aws/sensors/test_opensearch_serverless.py

@@ -0,0 +1,113 @@
+# 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 unittest import mock
+
+import pytest
+
+from airflow.exceptions import AirflowException, AirflowSkipException
+from airflow.providers.amazon.aws.hooks.opensearch_serverless import OpenSearchServerlessHook
+from airflow.providers.amazon.aws.sensors.opensearch_serverless import (
+    OpenSearchServerlessCollectionActiveSensor,
+)
+
+
+class TestOpenSearchServerlessCollectionActiveSensor:
+    def setup_method(self):
+        self.default_op_kwargs = dict(
+            task_id="test_sensor",
+            collection_id="knowledge_base_id",
+            poke_interval=5,
+            max_retries=1,
+        )
+        self.sensor = OpenSearchServerlessCollectionActiveSensor(**self.default_op_kwargs, aws_conn_id=None)
+
+    def test_base_aws_op_attributes(self):
+        op = OpenSearchServerlessCollectionActiveSensor(**self.default_op_kwargs)
+        assert op.hook.aws_conn_id == "aws_default"
+        assert op.hook._region_name is None
+        assert op.hook._verify is None
+        assert op.hook._config is None
+
+        op = OpenSearchServerlessCollectionActiveSensor(
+            **self.default_op_kwargs,
+            aws_conn_id="aws-test-custom-conn",
+            region_name="eu-west-1",
+            verify=False,
+            botocore_config={"read_timeout": 42},
+        )
+        assert op.hook.aws_conn_id == "aws-test-custom-conn"
+        assert op.hook._region_name == "eu-west-1"
+        assert op.hook._verify is False
+        assert op.hook._config is not None
+        assert op.hook._config.read_timeout == 42
+
+    @pytest.mark.parametrize(
+        "collection_name, collection_id, expected_pass",
+        [
+            pytest.param("name", "id", False, id="both_provided_fails"),
+            pytest.param("name", None, True, id="only_name_provided_passes"),
+            pytest.param(None, "id", True, id="only_id_provided_passes"),
+        ],
+    )
+    def test_name_and_id_combinations(self, collection_name, collection_id, expected_pass):
+        call_args = {
+            "task_id": "test_sensor",
+            "collection_name": collection_name,
+            "collection_id": collection_id,
+            "poke_interval": 5,
+            "max_retries": 1,
+        }
+        if expected_pass:
+            op = OpenSearchServerlessCollectionActiveSensor(**call_args)
+            assert op.collection_id == collection_id
+            assert op.collection_name == collection_name
+        if not expected_pass:
+            with pytest.raises(
+                AttributeError, match=OpenSearchServerlessCollectionActiveSensor.INVALID_ARGS_MESSAGE
+            ):
+                OpenSearchServerlessCollectionActiveSensor(**call_args)
+
+    @pytest.mark.parametrize("state", list(OpenSearchServerlessCollectionActiveSensor.SUCCESS_STATES))
+    @mock.patch.object(OpenSearchServerlessHook, "conn")
+    def test_poke_success_states(self, mock_conn, state):
+        mock_conn.batch_get_collection.return_value = {"collectionDetails": [{"status": state}]}
+        assert self.sensor.poke({}) is True
+
+    @pytest.mark.parametrize("state", list(OpenSearchServerlessCollectionActiveSensor.INTERMEDIATE_STATES))
+    @mock.patch.object(OpenSearchServerlessHook, "conn")
+    def test_poke_intermediate_states(self, mock_conn, state):
+        mock_conn.batch_get_collection.return_value = {"collectionDetails": [{"status": state}]}
+        assert self.sensor.poke({}) is False
+
+    @pytest.mark.parametrize(
+        "soft_fail, expected_exception",
+        [
+            pytest.param(False, AirflowException, id="not-soft-fail"),
+            pytest.param(True, AirflowSkipException, id="soft-fail"),
+        ],
+    )
+    @pytest.mark.parametrize("state", list(OpenSearchServerlessCollectionActiveSensor.FAILURE_STATES))
+    @mock.patch.object(OpenSearchServerlessHook, "conn")
+    def test_poke_failure_states(self, mock_conn, state, soft_fail, expected_exception):
+        mock_conn.batch_get_collection.return_value = {"collectionDetails": [{"status": state}]}
+        sensor = OpenSearchServerlessCollectionActiveSensor(
+            **self.default_op_kwargs, aws_conn_id=None, soft_fail=soft_fail
+        )
+        with pytest.raises(expected_exception, match=sensor.FAILURE_MESSAGE):
+            sensor.poke({})