diff --git a/providers/amazon/docs/operators/dynamodb.rst b/providers/amazon/docs/operators/dynamodb.rst index f802fcdb42389..633b311fcf0fd 100644 --- a/providers/amazon/docs/operators/dynamodb.rst +++ b/providers/amazon/docs/operators/dynamodb.rst @@ -45,6 +45,7 @@ Wait on Amazon DynamoDB item attribute value match Use the :class:`~airflow.providers.amazon.aws.sensors.dynamodb.DynamoDBValueSensor` to wait for the presence of a matching DynamoDB item's attribute/value pair. +This sensor can also be run in deferrable mode by setting ``deferrable`` param to ``True``. Wait for a Single Attribute Value Match: ---------------------------------------- diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index 4fa06f211212d..a45cdc0a0b7c0 100644 --- a/providers/amazon/provider.yaml +++ b/providers/amazon/provider.yaml @@ -821,6 +821,9 @@ triggers: - integration-name: Amazon Comprehend python-modules: - airflow.providers.amazon.aws.triggers.comprehend + - integration-name: Amazon DynamoDB + python-modules: + - airflow.providers.amazon.aws.triggers.dynamodb - integration-name: Amazon EC2 python-modules: - airflow.providers.amazon.aws.triggers.ec2 diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/dynamodb.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/dynamodb.py index c423605411008..321c639fb460a 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/dynamodb.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/dynamodb.py @@ -17,13 +17,17 @@ from __future__ import annotations from collections.abc import Iterable, Sequence +from datetime import timedelta from typing import TYPE_CHECKING, Any from botocore.exceptions import ClientError from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor +from airflow.providers.amazon.aws.triggers.dynamodb import DynamoDBValueSensorTrigger +from airflow.providers.amazon.aws.utils import validate_execute_complete_event from airflow.providers.amazon.aws.utils.mixins import aws_template_fields +from airflow.providers.common.compat.sdk import conf if TYPE_CHECKING: from airflow.sdk import Context @@ -44,6 +48,9 @@ class DynamoDBValueSensor(AwsBaseSensor[DynamoDBHook]): :param attribute_value: DynamoDB attribute value :param sort_key_name: (optional) DynamoDB sort key name :param sort_key_value: (optional) DynamoDB sort key value + :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) :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 @@ -76,6 +83,7 @@ def __init__( attribute_value: str | Iterable[str], sort_key_name: str | None = None, sort_key_value: str | None = None, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), **kwargs: Any, ): super().__init__(**kwargs) @@ -86,6 +94,37 @@ def __init__( self.attribute_value = attribute_value self.sort_key_name = sort_key_name self.sort_key_value = sort_key_value + self.deferrable = deferrable + + def execute(self, context: Context) -> Any: + if self.deferrable: + self.defer( + trigger=DynamoDBValueSensorTrigger( + table_name=self.table_name, + partition_key_name=self.partition_key_name, + partition_key_value=self.partition_key_value, + attribute_name=self.attribute_name, + attribute_value=self.attribute_value, + sort_key_name=self.sort_key_name, + sort_key_value=self.sort_key_value, + waiter_delay=int(self.poke_interval), + aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, + ), + method_name="execute_complete", + timeout=timedelta(seconds=self.timeout), + ) + else: + super().execute(context=context) + + def execute_complete(self, context: Context, event: dict | None = None) -> None: + validated_event = validate_execute_complete_event(event) + + if validated_event["status"] != "success": + raise RuntimeError(f"Trigger error: event is {validated_event}") + self.log.info("DynamoDB attribute value match found; sensor complete.") def poke(self, context: Context) -> bool: """Test DynamoDB item for matching attribute value.""" diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/dynamodb.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/dynamodb.py new file mode 100644 index 0000000000000..f1008b79f4df3 --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/dynamodb.py @@ -0,0 +1,168 @@ +# 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 + +import asyncio +from collections.abc import AsyncIterator, Iterable +from functools import cached_property +from typing import Any + +from boto3.dynamodb.types import TypeDeserializer +from botocore.exceptions import ClientError + +from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook +from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.triggers.base import BaseEventTrigger, TriggerEvent +else: + from airflow.triggers.base import ( # type: ignore + BaseTrigger as BaseEventTrigger, + TriggerEvent, + ) + + +class DynamoDBValueSensorTrigger(BaseEventTrigger): + """ + Asynchronously poll a DynamoDB item until the given attribute matches one of the expected values. + + The polling uses the low-level DynamoDB client (aiobotocore exposes clients only, not the + boto3 resource API used by ``DynamoDBHook``), so keys are sent as typed string attributes, + matching the string-typed key parameters of :class:`~airflow.providers.amazon.aws.sensors.dynamodb.DynamoDBValueSensor`. + + :param table_name: DynamoDB table name + :param partition_key_name: DynamoDB partition key name + :param partition_key_value: DynamoDB partition key value + :param attribute_name: DynamoDB attribute name + :param attribute_value: expected DynamoDB attribute value (or values, any of which matches) + :param sort_key_name: (optional) DynamoDB sort key name + :param sort_key_value: (optional) DynamoDB sort key value + :param waiter_delay: The time in seconds to wait between DynamoDB API calls + :param aws_conn_id: The Airflow connection used for AWS credentials. + :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 + """ + + def __init__( + self, + table_name: str, + partition_key_name: str, + partition_key_value: str, + attribute_name: str, + attribute_value: str | Iterable[str], + sort_key_name: str | None = None, + sort_key_value: str | None = None, + waiter_delay: int = 60, + aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, + ): + self.table_name = table_name + self.partition_key_name = partition_key_name + self.partition_key_value = partition_key_value + self.attribute_name = attribute_name + self.attribute_value = ( + [attribute_value] if isinstance(attribute_value, str) else list(attribute_value) + ) + self.sort_key_name = sort_key_name + self.sort_key_value = sort_key_value + self.waiter_delay = waiter_delay + self.aws_conn_id = aws_conn_id + self.region_name = region_name + self.verify = verify + self.botocore_config = botocore_config + + def serialize(self) -> tuple[str, dict[str, Any]]: + return ( + self.__class__.__module__ + "." + self.__class__.__qualname__, + { + "table_name": self.table_name, + "partition_key_name": self.partition_key_name, + "partition_key_value": self.partition_key_value, + "attribute_name": self.attribute_name, + "attribute_value": self.attribute_value, + "sort_key_name": self.sort_key_name, + "sort_key_value": self.sort_key_value, + "waiter_delay": self.waiter_delay, + "aws_conn_id": self.aws_conn_id, + "region_name": self.region_name, + "verify": self.verify, + "botocore_config": self.botocore_config, + }, + ) + + @cached_property + def hook(self) -> AwsBaseHook: + # DynamoDBHook is resource-based, but async connections are only available for + # clients, so the trigger talks to DynamoDB through a client-type hook instead. + return AwsBaseHook( + aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + config=self.botocore_config, + client_type="dynamodb", + ) + + @property + def key(self) -> dict[str, Any]: + key = {self.partition_key_name: {"S": self.partition_key_value}} + if self.sort_key_name and self.sort_key_value: + key[self.sort_key_name] = {"S": self.sort_key_value} + return key + + async def poke(self, client: Any) -> bool: + """Test the DynamoDB item for a matching attribute value, mirroring the sensor's poke.""" + self.log.info( + "Checking table %s for item with key %s, waiting for attribute %s to be one of %s", + self.table_name, + self.key, + self.attribute_name, + self.attribute_value, + ) + try: + response = await client.get_item(TableName=self.table_name, Key=self.key) + except ClientError as err: + # Same tolerance as the sensor's poke: log and keep trying until the task times out. + self.log.error( + "Couldn't get %s from table %s.\nError Code: %s\nError Message: %s", + self.key, + self.table_name, + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + return False + + try: + typed_attribute_value = response["Item"][self.attribute_name] + except KeyError: + return False + item_attribute_value = TypeDeserializer().deserialize(typed_attribute_value) + self.log.info("Got: %s = %s", self.attribute_name, item_attribute_value) + return item_attribute_value in self.attribute_value + + async def run(self) -> AsyncIterator[TriggerEvent]: + while True: + # This loop runs until the timeout set in the sensor's self.defer call is reached. + async with await self.hook.get_async_conn() as client: + if await self.poke(client=client): + yield TriggerEvent({"status": "success"}) + return + await asyncio.sleep(self.waiter_delay) diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py index a5a5c532e7e31..5736046d4d4ac 100644 --- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py +++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py @@ -921,6 +921,10 @@ def get_provider_info(): "integration-name": "Amazon Comprehend", "python-modules": ["airflow.providers.amazon.aws.triggers.comprehend"], }, + { + "integration-name": "Amazon DynamoDB", + "python-modules": ["airflow.providers.amazon.aws.triggers.dynamodb"], + }, { "integration-name": "Amazon EC2", "python-modules": ["airflow.providers.amazon.aws.triggers.ec2"], diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_dynamodb.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_dynamodb.py index f34acf3aac400..4fec505517ae0 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_dynamodb.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_dynamodb.py @@ -17,10 +17,13 @@ from __future__ import annotations +import pytest from moto import mock_aws from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook from airflow.providers.amazon.aws.sensors.dynamodb import DynamoDBValueSensor +from airflow.providers.amazon.aws.triggers.dynamodb import DynamoDBValueSensorTrigger +from airflow.providers.common.compat.sdk import TaskDeferred class TestDynamoDBValueSensor: @@ -118,6 +121,57 @@ def test_sensor_with_client_error(self): self.sensor_pk.partition_key_name = "no such key" assert self.sensor_pk.poke(None) is False + def test_sensor_deferrable(self): + sensor = DynamoDBValueSensor( + task_id="dynamodb_value_sensor_deferrable", + table_name=self.table_name, + partition_key_name=self.pk_name, + partition_key_value=self.pk_value, + attribute_name=self.attribute_name, + attribute_value=self.attribute_value, + sort_key_name=self.sk_name, + sort_key_value=self.sk_value, + deferrable=True, + ) + with pytest.raises(TaskDeferred) as defer: + sensor.execute({}) + trigger = defer.value.trigger + assert isinstance(trigger, DynamoDBValueSensorTrigger) + assert trigger.table_name == self.table_name + assert trigger.partition_key_name == self.pk_name + assert trigger.partition_key_value == self.pk_value + assert trigger.attribute_name == self.attribute_name + assert trigger.attribute_value == [self.attribute_value] + assert trigger.sort_key_name == self.sk_name + assert trigger.sort_key_value == self.sk_value + + def test_execute_complete(self): + sensor = DynamoDBValueSensor( + task_id="dynamodb_value_sensor_deferrable", + table_name=self.table_name, + partition_key_name=self.pk_name, + partition_key_value=self.pk_value, + attribute_name=self.attribute_name, + attribute_value=self.attribute_value, + deferrable=True, + ) + sensor.execute_complete(context={}, event={"status": "success"}) + + def test_fail_execute_complete(self): + sensor = DynamoDBValueSensor( + task_id="dynamodb_value_sensor_deferrable", + table_name=self.table_name, + partition_key_name=self.pk_name, + partition_key_value=self.pk_value, + attribute_name=self.attribute_name, + attribute_value=self.attribute_value, + deferrable=True, + ) + event = {"status": "failed"} + message = f"Trigger error: event is {event}" + with pytest.raises(RuntimeError, match=message): + sensor.execute_complete(context={}, event=event) + class TestDynamoDBMultipleValuesSensor: def setup_method(self): diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_dynamodb.py b/providers/amazon/tests/unit/amazon/aws/triggers/test_dynamodb.py new file mode 100644 index 0000000000000..afcd199138482 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_dynamodb.py @@ -0,0 +1,175 @@ +# 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 +from unittest.mock import AsyncMock + +import pytest +from botocore.exceptions import ClientError + +from airflow.providers.amazon.aws.triggers.dynamodb import DynamoDBValueSensorTrigger +from airflow.triggers.base import TriggerEvent + +TEST_TABLE_NAME = "test-table" +TEST_PK_NAME = "PK" +TEST_PK_VALUE = "PKTest" +TEST_SK_NAME = "SK" +TEST_SK_VALUE = "SKTest" +TEST_ATTRIBUTE_NAME = "Foo" +TEST_ATTRIBUTE_VALUE = "Bar" +TEST_WAITER_DELAY = 1 +TEST_CONN_ID = "test-conn-id" +TEST_REGION_NAME = "eu-central-1" + + +def _build_trigger(**kwargs) -> DynamoDBValueSensorTrigger: + params = { + "table_name": TEST_TABLE_NAME, + "partition_key_name": TEST_PK_NAME, + "partition_key_value": TEST_PK_VALUE, + "attribute_name": TEST_ATTRIBUTE_NAME, + "attribute_value": TEST_ATTRIBUTE_VALUE, + "waiter_delay": TEST_WAITER_DELAY, + "aws_conn_id": TEST_CONN_ID, + "region_name": TEST_REGION_NAME, + **kwargs, + } + return DynamoDBValueSensorTrigger(**params) + + +class TestDynamoDBValueSensorTrigger: + def test_serialize(self): + trigger = _build_trigger(sort_key_name=TEST_SK_NAME, sort_key_value=TEST_SK_VALUE) + + class_path, args = trigger.serialize() + assert class_path == "airflow.providers.amazon.aws.triggers.dynamodb.DynamoDBValueSensorTrigger" + assert args["table_name"] == TEST_TABLE_NAME + assert args["partition_key_name"] == TEST_PK_NAME + assert args["partition_key_value"] == TEST_PK_VALUE + assert args["attribute_name"] == TEST_ATTRIBUTE_NAME + # a single string value is normalized to a list so serialization round-trips + assert args["attribute_value"] == [TEST_ATTRIBUTE_VALUE] + assert args["sort_key_name"] == TEST_SK_NAME + assert args["sort_key_value"] == TEST_SK_VALUE + assert args["waiter_delay"] == TEST_WAITER_DELAY + assert args["aws_conn_id"] == TEST_CONN_ID + assert args["region_name"] == TEST_REGION_NAME + assert args["verify"] is None + assert args["botocore_config"] is None + + def test_serialize_generic_hook_params(self): + trigger = _build_trigger(verify=False, botocore_config={"read_timeout": 99}) + _, args = trigger.serialize() + assert args["verify"] is False + assert args["botocore_config"] == {"read_timeout": 99} + + hook = trigger.hook + assert hook.client_type == "dynamodb" + assert hook.aws_conn_id == TEST_CONN_ID + assert hook._region_name == TEST_REGION_NAME + assert hook._verify is False + assert hook._config.read_timeout == 99 + + def test_key_with_and_without_sort_key(self): + trigger = _build_trigger() + assert trigger.key == {TEST_PK_NAME: {"S": TEST_PK_VALUE}} + + trigger = _build_trigger(sort_key_name=TEST_SK_NAME, sort_key_value=TEST_SK_VALUE) + assert trigger.key == { + TEST_PK_NAME: {"S": TEST_PK_VALUE}, + TEST_SK_NAME: {"S": TEST_SK_VALUE}, + } + + @pytest.mark.asyncio + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_async_conn") + async def test_run_success_when_value_matches(self, mock_async_conn): + client = AsyncMock() + mock_async_conn.return_value.__aenter__.return_value = client + client.get_item.return_value = { + "Item": { + TEST_PK_NAME: {"S": TEST_PK_VALUE}, + TEST_ATTRIBUTE_NAME: {"S": TEST_ATTRIBUTE_VALUE}, + } + } + + trigger = _build_trigger() + generator = trigger.run() + response = await generator.asend(None) + + client.get_item.assert_called_once_with( + TableName=TEST_TABLE_NAME, Key={TEST_PK_NAME: {"S": TEST_PK_VALUE}} + ) + assert response == TriggerEvent({"status": "success"}) + + @pytest.mark.asyncio + @mock.patch("asyncio.sleep") + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_async_conn") + async def test_run_keeps_polling_until_value_matches(self, mock_async_conn, mock_sleep): + mock_sleep.return_value = True + client = AsyncMock() + mock_async_conn.return_value.__aenter__.return_value = client + client.get_item.side_effect = [ + {}, # item does not exist yet + {"Item": {TEST_PK_NAME: {"S": TEST_PK_VALUE}, TEST_ATTRIBUTE_NAME: {"S": "wrong"}}}, + {"Item": {TEST_PK_NAME: {"S": TEST_PK_VALUE}, TEST_ATTRIBUTE_NAME: {"S": TEST_ATTRIBUTE_VALUE}}}, + ] + + trigger = _build_trigger() + generator = trigger.run() + response = await generator.asend(None) + + assert client.get_item.call_count == 3 + assert response == TriggerEvent({"status": "success"}) + + @pytest.mark.asyncio + @mock.patch("asyncio.sleep") + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_async_conn") + async def test_run_tolerates_client_error(self, mock_async_conn, mock_sleep): + # Same tolerance as the sensor's poke: a ClientError means "not there yet", not failure. + mock_sleep.return_value = True + client = AsyncMock() + mock_async_conn.return_value.__aenter__.return_value = client + client.get_item.side_effect = [ + ClientError( + error_response={"Error": {"Code": "ResourceNotFoundException", "Message": "no table"}}, + operation_name="GetItem", + ), + {"Item": {TEST_PK_NAME: {"S": TEST_PK_VALUE}, TEST_ATTRIBUTE_NAME: {"S": TEST_ATTRIBUTE_VALUE}}}, + ] + + trigger = _build_trigger() + generator = trigger.run() + response = await generator.asend(None) + + assert client.get_item.call_count == 2 + assert response == TriggerEvent({"status": "success"}) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_async_conn") + async def test_run_success_with_multiple_expected_values(self, mock_async_conn): + client = AsyncMock() + mock_async_conn.return_value.__aenter__.return_value = client + client.get_item.return_value = { + "Item": {TEST_PK_NAME: {"S": TEST_PK_VALUE}, TEST_ATTRIBUTE_NAME: {"S": "Bar2"}} + } + + trigger = _build_trigger(attribute_value=["Bar1", "Bar2", "Bar3"]) + generator = trigger.run() + response = await generator.asend(None) + + assert response == TriggerEvent({"status": "success"})