apache-airflow-providers-amazon 9.2.0rc2__py3-none-any.whl → 9.3.0__py3-none-any.whl

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

@@ -145,13 +145,13 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
         return execution_arn
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Trigger error: event is {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Trigger error: event is {validated_event}")
 
         self.log.info("State Machine execution completed successfully")
-        return event["execution_arn"]
+        return validated_event["execution_arn"]
 
 
 class StepFunctionGetExecutionOutputOperator(AwsBaseOperator[StepFunctionHook]):

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

@@ -95,7 +95,7 @@ class EC2InstanceStateSensor(BaseSensorOperator):
         return instance_state == self.target_state
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Error: {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Error: {validated_event}")

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

@@ -346,10 +346,10 @@ class EmrContainerSensor(BaseSensorOperator):
             )
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        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.")
 
@@ -535,10 +535,10 @@ class EmrJobFlowSensor(EmrBaseSensor):
             )
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        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.")
 
 
@@ -664,9 +664,9 @@ class EmrStepSensor(EmrBaseSensor):
             )
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        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 %s completed.", self.job_flow_id)

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

@@ -174,14 +174,14 @@ class GlueDataQualityRuleSetEvaluationRunSensor(AwsBaseSensor[GlueDataQualityHoo
             super().execute(context=context)
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            message = f"Error: AWS Glue data quality ruleset evaluation run: {event}"
+        if validated_event["status"] != "success":
+            message = f"Error: AWS Glue data quality ruleset evaluation run: {validated_event}"
             raise AirflowException(message)
 
         self.hook.validate_evaluation_run_results(
-            evaluation_run_id=event["evaluation_run_id"],
+            evaluation_run_id=validated_event["evaluation_run_id"],
             show_results=self.show_results,
             verify_result_status=self.verify_result_status,
         )
@@ -295,10 +295,10 @@ class GlueDataQualityRuleRecommendationRunSensor(AwsBaseSensor[GlueDataQualityHo
             super().execute(context=context)
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            message = f"Error: AWS Glue data quality recommendation run: {event}"
+        if validated_event["status"] != "success":
+            message = f"Error: AWS Glue data quality recommendation run: {validated_event}"
             raise AirflowException(message)
 
         if self.show_results:

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

@@ -123,8 +123,8 @@ class GlueCatalogPartitionSensor(AwsBaseSensor[GlueCatalogHook]):
         return self.hook.check_for_partition(self.database_name, self.table_name, self.expression)
 
     def execute_complete(self, context: Context, event: dict | None = None) -> None:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Trigger error: event is {event}")
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Trigger error: event is {validated_event}")
         self.log.info("Partition exists in the Glue Catalog")

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

@@ -88,11 +88,11 @@ class RedshiftClusterSensor(BaseSensorOperator):
             )
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        status = event["status"]
+        status = validated_event["status"]
         if status == "error":
-            raise AirflowException(f"{event['status']}: {event['message']}")
+            raise AirflowException(f"{validated_event['status']}: {validated_event['message']}")
         elif status == "success":
             self.log.info("%s completed successfully.", self.task_id)
             self.log.info("Cluster Identifier %s is in %s state", self.cluster_identifier, self.target_status)

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

@@ -30,10 +30,11 @@ from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
 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.amazon.aws.utils.sqs import MessageFilteringType, process_response
+from airflow.providers.amazon.aws.utils.sqs import process_response
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
+    from airflow.providers.amazon.aws.utils.sqs import MessageFilteringType
     from airflow.utils.context import Context
 
 
@@ -156,11 +157,11 @@ class SqsSensor(AwsBaseSensor[SqsHook]):
             super().execute(context=context)
 
     def execute_complete(self, context: Context, event: dict | None = None) -> None:
-        event = validate_execute_complete_event(event)
+        validated_event = validate_execute_complete_event(event)
 
-        if event["status"] != "success":
-            raise AirflowException(f"Trigger error: event is {event}")
-        context["ti"].xcom_push(key="messages", value=event["message_batch"])
+        if validated_event["status"] != "success":
+            raise AirflowException(f"Trigger error: event is {validated_event}")
+        context["ti"].xcom_push(key="messages", value=validated_event["message_batch"])
 
     def poll_sqs(self, sqs_conn: BaseAwsConnection) -> Collection:
         """

airflow/providers/amazon/aws/transfers/google_api_to_s3.py

@@ -30,7 +30,10 @@ from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.google.common.hooks.discovery_api import GoogleDiscoveryApiHook
 
 if TYPE_CHECKING:
-    from airflow.models import TaskInstance
+    try:
+        from airflow.sdk.types import RuntimeTaskInstanceProtocol
+    except ImportError:
+        from airflow.models import TaskInstance as RuntimeTaskInstanceProtocol  # type: ignore[assignment]
     from airflow.utils.context import Context
 
 
@@ -174,7 +177,7 @@ class GoogleApiToS3Operator(BaseOperator):
             replace=self.s3_overwrite,
         )
 
-    def _update_google_api_endpoint_params_via_xcom(self, task_instance: TaskInstance) -> None:
+    def _update_google_api_endpoint_params_via_xcom(self, task_instance: RuntimeTaskInstanceProtocol) -> None:
         if self.google_api_endpoint_params_via_xcom:
             google_api_endpoint_params = task_instance.xcom_pull(
                 task_ids=self.google_api_endpoint_params_via_xcom_task_ids,
@@ -182,7 +185,9 @@ class GoogleApiToS3Operator(BaseOperator):
             )
             self.google_api_endpoint_params.update(google_api_endpoint_params)
 
-    def _expose_google_api_response_via_xcom(self, task_instance: TaskInstance, data: dict) -> None:
+    def _expose_google_api_response_via_xcom(
+        self, task_instance: RuntimeTaskInstanceProtocol, data: dict
+    ) -> None:
         if sys.getsizeof(data) < MAX_XCOM_SIZE:
             task_instance.xcom_push(key=self.google_api_response_via_xcom or XCOM_RETURN_KEY, value=data)
         else:

airflow/providers/amazon/aws/triggers/README.md

@@ -126,7 +126,7 @@ Here, we are calling the `get_waiter` function defined in `base_aws.py` which ta
 }
 ```
 
-For more information about writing custom waiter, see the [README.md](https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/waiters/README.md) for custom waiters.
+For more information about writing custom waiter, see the [README.md](https://github.com/apache/airflow/blob/main/providers/amazon/aws/src/airflow/providers/amazon/aws/waiters/README.md) for custom waiters.
 
 In some cases, a built-in or custom waiter may not be able to solve the problem. In such cases, the asynchronous method used to poll the boto3 API would need to be defined in the hook of the service being used. This method is essentially the same as the synchronous version of the method, except that it will use the aiobotocore client, and will be awaited. For the Redshift example, the async `describe_clusters` method would look as follows:
 

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

@@ -53,7 +53,8 @@ class OpenSearchServerlessCollectionActiveTrigger(AwsBaseWaiterTrigger):
         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]},
+            # waiter_args is a dict[str, Any], allow a possible list of None (it is caught above)
+            waiter_args={"ids": [collection_id]} if collection_id else {"names": [collection_name]},  # type: ignore[list-item]
             failure_message="OpenSearch Serverless Collection creation failed.",
             status_message="Status of OpenSearch Serverless Collection is",
             status_queries=["status"],

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

@@ -22,11 +22,12 @@ from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
-from airflow.providers.amazon.aws.utils.sqs import MessageFilteringType, process_response
+from airflow.providers.amazon.aws.utils.sqs import process_response
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
+    from airflow.providers.amazon.aws.utils.sqs import MessageFilteringType
 
 
 class SqsSensorTrigger(BaseTrigger):

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

@@ -14,20 +14,22 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+#
 from __future__ import annotations
 
 import json
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import jsonpath_ng
 import jsonpath_ng.ext
-from typing_extensions import Literal
 
-log = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from typing import Literal
 
+    MessageFilteringType = Literal["literal", "jsonpath", "jsonpath-ext"]
 
-MessageFilteringType = Literal["literal", "jsonpath", "jsonpath-ext"]
+log = logging.getLogger(__name__)
 
 
 def process_response(

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

@@ -36,6 +36,18 @@
                     "argument": "Replications[0].Status",
                     "expected": "stopped",
                     "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "Replications[0].Status",
+                    "expected": "created",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "Replications[0].ProvisionData.ProvisionState",
+                    "expected": "deprovisioned",
+                    "state": "success"
                 }
             ]
         },