diff --git a/airflow/providers/amazon/aws/sensors/dynamodb.py b/airflow/providers/amazon/aws/sensors/dynamodb.py index dbb7f973041e..ead8c123a621 100644 --- a/airflow/providers/amazon/aws/sensors/dynamodb.py +++ b/airflow/providers/amazon/aws/sensors/dynamodb.py @@ -18,6 +18,8 @@ from typing import TYPE_CHECKING, Any, Iterable, Sequence +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.utils.mixins import aws_template_fields @@ -102,14 +104,26 @@ def poke(self, context: Context) -> bool: table = self.hook.conn.Table(self.table_name) self.log.info("Table: %s", table) self.log.info("Key: %s", key) - response = table.get_item(Key=key) + try: - item_attribute_value = response["Item"][self.attribute_name] - self.log.info("Response: %s", response) - self.log.info("Want: %s = %s", self.attribute_name, self.attribute_value) - self.log.info("Got: {response['Item'][self.attribute_name]} = %s", item_attribute_value) - return item_attribute_value in ( - [self.attribute_value] if isinstance(self.attribute_value, str) else self.attribute_value + response = table.get_item(Key=key) + except ClientError as err: + self.log.error( + "Couldn't get %s from table %s.\nError Code: %s\nError Message: %s", + key, + self.table_name, + err.response["Error"]["Code"], + err.response["Error"]["Message"], ) - except KeyError: return False + else: + try: + item_attribute_value = response["Item"][self.attribute_name] + self.log.info("Response: %s", response) + self.log.info("Want: %s = %s", self.attribute_name, self.attribute_value) + self.log.info("Got: {response['Item'][self.attribute_name]} = %s", item_attribute_value) + return item_attribute_value in ( + [self.attribute_value] if isinstance(self.attribute_value, str) else self.attribute_value + ) + except KeyError: + return False diff --git a/tests/providers/amazon/aws/sensors/test_dynamodb.py b/tests/providers/amazon/aws/sensors/test_dynamodb.py index d8b31b48c5e3..93ca01d26275 100644 --- a/tests/providers/amazon/aws/sensors/test_dynamodb.py +++ b/tests/providers/amazon/aws/sensors/test_dynamodb.py @@ -104,6 +104,23 @@ def test_sensor_with_pk_and_sk(self): assert self.sensor_pk_sk.poke(None) + @mock_aws + def test_sensor_with_client_error(self): + hook = DynamoDBHook(table_name=self.table_name, table_keys=[self.pk_name]) + + hook.conn.create_table( + TableName=self.table_name, + KeySchema=[{"AttributeName": self.pk_name, "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": self.pk_name, "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 10, "WriteCapacityUnits": 10}, + ) + + items = [{self.pk_name: self.pk_value, self.attribute_name: self.attribute_value}] + hook.write_batch_data(items) + + self.sensor_pk.partition_key_name = "no such key" + assert self.sensor_pk.poke(None) is False + class TestDynamoDBMultipleValuesSensor: def setup_method(self):