apache-airflow-providers-amazon 9.2.0rc2__py3-none-any.whl → 9.4.0rc1__py3-none-any.whl

airflow/providers/amazon/aws/hooks/base_aws.py

@@ -58,7 +58,6 @@ from airflow.providers.amazon.aws.utils.suppress import return_on_error
 from airflow.providers_manager import ProvidersManager
 from airflow.utils.helpers import exactly_one
 from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.log.secrets_masker import mask_secret
 
 BaseAwsConnection = TypeVar("BaseAwsConnection", bound=Union[boto3.client, boto3.resource])
 
@@ -68,6 +67,12 @@ if TYPE_CHECKING:
     from botocore.credentials import ReadOnlyCredentials
 
     from airflow.models.connection import Connection
+    from airflow.sdk.execution_time.secrets_masker import mask_secret
+else:
+    try:
+        from airflow.sdk.execution_time.secrets_masker import mask_secret
+    except ImportError:
+        from airflow.utils.log.secrets_masker import mask_secret
 
 _loader = botocore.loaders.Loader()
 """

airflow/providers/amazon/aws/hooks/batch_client.py

@@ -30,7 +30,7 @@ from __future__ import annotations
 import itertools
 import random
 import time
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, Protocol, runtime_checkable
 
 import botocore.client
 import botocore.exceptions
@@ -38,7 +38,6 @@ import botocore.waiter
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.typing_compat import Protocol, runtime_checkable
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher

airflow/providers/amazon/aws/hooks/ecr.py

@@ -23,11 +23,17 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.utils.log.secrets_masker import mask_secret
 
 if TYPE_CHECKING:
     from datetime import datetime
 
+    from airflow.sdk.execution_time.secrets_masker import mask_secret
+else:
+    try:
+        from airflow.sdk.execution_time.secrets_masker import mask_secret
+    except ImportError:
+        from airflow.utils.log.secrets_masker import mask_secret
+
 logger = logging.getLogger(__name__)
 
 

airflow/providers/amazon/aws/hooks/ecs.py

@@ -17,12 +17,11 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 from airflow.providers.amazon.aws.exceptions import EcsOperatorError, EcsTaskFailToStart
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
 from airflow.providers.amazon.aws.utils import _StringCompareEnum
-from airflow.typing_compat import Protocol, runtime_checkable
 
 if TYPE_CHECKING:
     from botocore.waiter import Waiter

airflow/providers/amazon/aws/hooks/eks.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 
 import base64
 import json
+import os
 import sys
 import tempfile
 from collections.abc import Generator
@@ -32,6 +33,7 @@ from botocore.exceptions import ClientError
 from botocore.signers import RequestSigner
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.hooks.sts import StsHook
 from airflow.utils import yaml
 from airflow.utils.json import AirflowJsonEncoder
 
@@ -612,9 +614,14 @@ class EksHook(AwsBaseHook):
     def fetch_access_token_for_cluster(self, eks_cluster_name: str) -> str:
         session = self.get_session()
         service_id = self.conn.meta.service_model.service_id
-        sts_url = (
-            f"https://sts.{session.region_name}.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15"
-        )
+        # This env variable is required so that we get a regionalized endpoint for STS in regions that
+        # otherwise default to global endpoints. The mechanism below to generate the token is very picky that
+        # the endpoint is regional.
+        os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional"
+        try:
+            sts_url = f"{StsHook(region_name=session.region_name).conn_client_meta.endpoint_url}/?Action=GetCallerIdentity&Version=2011-06-15"
+        finally:
+            del os.environ["AWS_STS_REGIONAL_ENDPOINTS"]
 
         signer = RequestSigner(
             service_id=service_id,

airflow/providers/amazon/aws/hooks/emr.py

@@ -22,7 +22,9 @@ import time
 import warnings
 from typing import Any
 
+import tenacity
 from botocore.exceptions import ClientError
+from tenacity import retry_if_exception, stop_after_attempt, wait_fixed
 
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -311,6 +313,15 @@ class EmrServerlessHook(AwsBaseHook):
         return count
 
 
+def is_connection_being_updated_exception(exception: BaseException) -> bool:
+    return (
+        isinstance(exception, ClientError)
+        and exception.response["Error"]["Code"] == "ValidationException"
+        and "is not reachable as its connection is currently being updated"
+        in exception.response["Error"]["Message"]
+    )
+
+
 class EmrContainerHook(AwsBaseHook):
     """
     Interact with Amazon EMR Containers (Amazon EMR on EKS).
@@ -348,6 +359,15 @@ class EmrContainerHook(AwsBaseHook):
         super().__init__(client_type="emr-containers", *args, **kwargs)  # type: ignore
         self.virtual_cluster_id = virtual_cluster_id
 
+    # Retry this method when the ``create_virtual_cluster`` raises
+    # "Cluster XXX is not reachable as its connection is currently being updated".
+    # Even though the EKS cluster status is ``ACTIVE``, ``create_virtual_cluster`` can raise this error.
+    # Retrying is the only option. Retry up to 3 minutes
+    @tenacity.retry(
+        retry=retry_if_exception(is_connection_being_updated_exception),
+        stop=stop_after_attempt(12),
+        wait=wait_fixed(15),
+    )
     def create_emr_on_eks_cluster(
         self,
         virtual_cluster_name: str,

airflow/providers/amazon/aws/hooks/mwaa.py

@@ -0,0 +1,85 @@
+# 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.
+"""This module contains AWS MWAA hook."""
+
+from __future__ import annotations
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class MwaaHook(AwsBaseHook):
+    """
+    Interact with AWS Manager Workflows for Apache Airflow.
+
+    Provide thin wrapper around :external+boto3:py:class:`boto3.client("mwaa") <MWAA.Client>`
+
+    Additional arguments (such as ``aws_conn_id``) may be specified and
+    are passed down to the underlying AwsBaseHook.
+
+    .. seealso::
+        - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        kwargs["client_type"] = "mwaa"
+        super().__init__(*args, **kwargs)
+
+    def invoke_rest_api(
+        self,
+        env_name: str,
+        path: str,
+        method: str,
+        body: dict | None = None,
+        query_params: dict | None = None,
+    ) -> dict:
+        """
+        Invoke the REST API on the Airflow webserver with the specified inputs.
+
+        .. seealso::
+            - :external+boto3:py:meth:`MWAA.Client.invoke_rest_api`
+
+        :param env_name: name of the MWAA environment
+        :param path: Apache Airflow REST API endpoint path to be called
+        :param method: HTTP method used for making Airflow REST API calls
+        :param body: Request body for the Apache Airflow REST API call
+        :param query_params: Query parameters to be included in the Apache Airflow REST API call
+        """
+        body = body or {}
+        api_kwargs = {
+            "Name": env_name,
+            "Path": path,
+            "Method": method,
+            # Filter out keys with None values because Airflow REST API doesn't accept requests otherwise
+            "Body": {k: v for k, v in body.items() if v is not None},
+            "QueryParameters": query_params if query_params else {},
+        }
+        try:
+            result = self.conn.invoke_rest_api(**api_kwargs)
+            # ResponseMetadata is removed because it contains data that is either very unlikely to be useful
+            # in XComs and logs, or redundant given the data already included in the response
+            result.pop("ResponseMetadata", None)
+            return result
+        except ClientError as e:
+            to_log = e.response
+            # ResponseMetadata and Error are removed because they contain data that is either very unlikely to
+            # be useful in XComs and logs, or redundant given the data already included in the response
+            to_log.pop("ResponseMetadata", None)
+            to_log.pop("Error", None)
+            self.log.error(to_log)
+            raise e

airflow/providers/amazon/aws/hooks/sqs.py

@@ -59,6 +59,7 @@ class SqsHook(AwsBaseHook):
         delay_seconds: int = 0,
         message_attributes: dict | None = None,
         message_group_id: str | None = None,
+        message_deduplication_id: str | None = None,
     ) -> dict:
         """
         Send message to the queue.
@@ -71,6 +72,7 @@ class SqsHook(AwsBaseHook):
         :param delay_seconds: seconds to delay the message
         :param message_attributes: additional attributes for the message (default: None)
         :param message_group_id: This applies only to FIFO (first-in-first-out) queues. (default: None)
+        :param message_deduplication_id: This applies only to FIFO (first-in-first-out) queues.
         :return: dict with the information about the message sent
         """
         params = {
@@ -81,5 +83,7 @@ class SqsHook(AwsBaseHook):
         }
         if message_group_id:
             params["MessageGroupId"] = message_group_id
+        if message_deduplication_id:
+            params["MessageDeduplicationId"] = message_deduplication_id
 
         return self.get_conn().send_message(**params)

airflow/providers/amazon/aws/hooks/ssm.py

@@ -17,10 +17,19 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.utils.log.secrets_masker import mask_secret
 from airflow.utils.types import NOTSET, ArgNotSet
 
+if TYPE_CHECKING:
+    from airflow.sdk.execution_time.secrets_masker import mask_secret
+else:
+    try:
+        from airflow.sdk.execution_time.secrets_masker import mask_secret
+    except ImportError:
+        from airflow.utils.log.secrets_masker import mask_secret
+
 
 class SsmHook(AwsBaseHook):
     """

airflow/providers/amazon/aws/links/comprehend.py

@@ -0,0 +1,41 @@
+# 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 airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+
+
+class ComprehendPiiEntitiesDetectionLink(BaseAwsLink):
+    """Helper class for constructing Amazon Comprehend PII Detection console link."""
+
+    name = "PII Detection Job"
+    key = "comprehend_pii_detection"
+    format_str = (
+        BASE_AWS_CONSOLE_LINK
+        + "/comprehend/home?region={region_name}#"
+        + "/analysis-job-details/pii/{job_id}"
+    )
+
+
+class ComprehendDocumentClassifierLink(BaseAwsLink):
+    """Helper class for constructing Amazon Comprehend Document Classifier console link."""
+
+    name = "Document Classifier"
+    key = "comprehend_document_classifier"
+    format_str = (
+        BASE_AWS_CONSOLE_LINK + "/comprehend/home?region={region_name}#" + "classifier-version-details/{arn}"
+    )

airflow/providers/amazon/aws/links/datasync.py

@@ -0,0 +1,37 @@
+# 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 airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+
+
+class DataSyncTaskLink(BaseAwsLink):
+    """Helper class for constructing AWS DataSync Task console link."""
+
+    name = "DataSync Task"
+    key = "datasync_task"
+    format_str = BASE_AWS_CONSOLE_LINK + "/datasync/home?region={region_name}#" + "/tasks/{task_id}"
+
+
+class DataSyncTaskExecutionLink(BaseAwsLink):
+    """Helper class for constructing AWS DataSync TaskExecution console link."""
+
+    name = "DataSync Task Execution"
+    key = "datasync_task_execution"
+    format_str = (
+        BASE_AWS_CONSOLE_LINK + "/datasync/home?region={region_name}#/history/{task_id}/{task_execution_id}"
+    )

airflow/providers/amazon/aws/links/ec2.py

@@ -0,0 +1,46 @@
+# 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 airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+
+
+class EC2InstanceLink(BaseAwsLink):
+    """Helper class for constructing Amazon EC2 instance links."""
+
+    name = "Instance"
+    key = "_instance_id"
+    format_str = (
+        BASE_AWS_CONSOLE_LINK + "/ec2/home?region={region_name}#InstanceDetails:instanceId={instance_id}"
+    )
+
+
+class EC2InstanceDashboardLink(BaseAwsLink):
+    """
+    Helper class for constructing Amazon EC2 console links.
+
+    This is useful for displaying the list of EC2 instances, rather
+    than a single instance.
+    """
+
+    name = "EC2 Instances"
+    key = "_instance_dashboard"
+    format_str = BASE_AWS_CONSOLE_LINK + "/ec2/home?region={region_name}#Instances:instanceId=:{instance_ids}"
+
+    @staticmethod
+    def format_instance_id_filter(instance_ids: list[str]) -> str:
+        return ",:".join(instance_ids)

airflow/providers/amazon/aws/links/sagemaker.py

@@ -0,0 +1,27 @@
+# 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 airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+
+
+class SageMakerTransformJobLink(BaseAwsLink):
+    """Helper class for constructing AWS Transform Run Details Link."""
+
+    name = "Amazon SageMaker Transform Job Details"
+    key = "sagemaker_transform_job_details"
+    format_str = BASE_AWS_CONSOLE_LINK + "/sagemaker/home?region={region_name}#/transform-jobs/{job_name}"

airflow/providers/amazon/aws/operators/athena.py

@@ -177,14 +177,16 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]):
         return self.query_execution_id
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Error while waiting for operation on cluster to complete: {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(
+                f"Error while waiting for operation on cluster to complete: {validated_event}"
+            )
 
         # Save query_execution_id to be later used by listeners
-        self.query_execution_id = event["value"]
-        return event["value"]
+        self.query_execution_id = validated_event["value"]
+        return validated_event["value"]
 
     def on_kill(self) -> None:
         """Cancel the submitted Amazon Athena query."""

airflow/providers/amazon/aws/operators/batch.py

@@ -95,6 +95,7 @@ class BatchOperator(BaseOperator):
         If it is an array job, only the logs of the first task will be printed.
     :param awslogs_fetch_interval: The interval with which cloudwatch logs are to be fetched, 30 sec.
     :param poll_interval: (Deferrable mode only) Time in seconds to wait between polling.
+    :param submit_job_timeout: Execution timeout in seconds for submitted batch job.
 
     .. note::
         Any custom waiters must return a waiter for these calls:
@@ -184,6 +185,7 @@ class BatchOperator(BaseOperator):
         poll_interval: int = 30,
         awslogs_enabled: bool = False,
         awslogs_fetch_interval: timedelta = timedelta(seconds=30),
+        submit_job_timeout: int | None = None,
         **kwargs,
     ) -> None:
         BaseOperator.__init__(self, **kwargs)
@@ -208,6 +210,7 @@ class BatchOperator(BaseOperator):
         self.poll_interval = poll_interval
         self.awslogs_enabled = awslogs_enabled
         self.awslogs_fetch_interval = awslogs_fetch_interval
+        self.submit_job_timeout = submit_job_timeout
 
         # params for hook
         self.max_retries = max_retries
@@ -264,13 +267,13 @@ class BatchOperator(BaseOperator):
         return self.job_id
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Error while running job: {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Error while running job: {validated_event}")
 
         self.log.info("Job completed.")
-        return event["job_id"]
+        return validated_event["job_id"]
 
     def on_kill(self):
         response = self.hook.client.terminate_job(jobId=self.job_id, reason="Task killed by the user")
@@ -315,6 +318,9 @@ class BatchOperator(BaseOperator):
             "schedulingPriorityOverride": self.scheduling_priority_override,
         }
 
+        if self.submit_job_timeout:
+            args["timeout"] = {"attemptDurationSeconds": self.submit_job_timeout}
+
         try:
             response = self.hook.client.submit_job(**trim_none_values(args))
         except Exception as e:
@@ -534,8 +540,10 @@ class BatchCreateComputeEnvironmentOperator(BaseOperator):
         return arn
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Error while waiting for the compute environment to be ready: {event}")
-        return event["value"]
+        if validated_event["status"] != "success":
+            raise AirflowException(
+                f"Error while waiting for the compute environment to be ready: {validated_event}"
+            )
+        return validated_event["value"]

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

@@ -198,13 +198,13 @@ class BedrockCustomizeModelOperator(AwsBaseOperator[BedrockHook]):
         self.valid_action_if_job_exists: set[str] = {"timestamp", "fail"}
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Error while running job: {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Error while running job: {validated_event}")
 
         self.log.info("Bedrock model customization job `%s` complete.", self.job_name)
-        return self.hook.conn.get_model_customization_job(jobIdentifier=event["job_name"])["jobArn"]
+        return self.hook.conn.get_model_customization_job(jobIdentifier=validated_event["job_name"])["jobArn"]
 
     def execute(self, context: Context) -> dict:
         response = {}
@@ -353,13 +353,15 @@ class BedrockCreateProvisionedModelThroughputOperator(AwsBaseOperator[BedrockHoo
         return provisioned_model_id
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Error while running job: {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Error while running job: {validated_event}")
 
-        self.log.info("Bedrock provisioned throughput job `%s` complete.", event["provisioned_model_id"])
-        return event["provisioned_model_id"]
+        self.log.info(
+            "Bedrock provisioned throughput job `%s` complete.", validated_event["provisioned_model_id"]
+        )
+        return validated_event["provisioned_model_id"]
 
 
 class BedrockCreateKnowledgeBaseOperator(AwsBaseOperator[BedrockAgentHook]):
@@ -449,13 +451,13 @@ class BedrockCreateKnowledgeBaseOperator(AwsBaseOperator[BedrockAgentHook]):
         self.deferrable = deferrable
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Error while running job: {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Error while running job: {validated_event}")
 
         self.log.info("Bedrock knowledge base creation job `%s` complete.", self.name)
-        return event["knowledge_base_id"]
+        return validated_event["knowledge_base_id"]
 
     def execute(self, context: Context) -> str:
         def _create_kb():
@@ -634,14 +636,14 @@ class BedrockIngestDataOperator(AwsBaseOperator[BedrockAgentHook]):
         self.deferrable = deferrable
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Error while running ingestion job: {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Error while running ingestion job: {validated_event}")
 
-        self.log.info("Bedrock ingestion job `%s` complete.", event["ingestion_job_id"])
+        self.log.info("Bedrock ingestion job `%s` complete.", validated_event["ingestion_job_id"])
 
-        return event["ingestion_job_id"]
+        return validated_event["ingestion_job_id"]
 
     def execute(self, context: Context) -> str:
         ingestion_job_id = self.hook.conn.start_ingestion_job(