apache-airflow-providers-amazon 9.5.0rc3__py3-none-any.whl → 9.6.0rc1__py3-none-any.whl

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

@@ -17,36 +17,30 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from functools import cached_property
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
-from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class RdsBaseSensor(BaseSensorOperator):
+class RdsBaseSensor(AwsBaseSensor[RdsHook]):
     """Base operator that implements common functions for all sensors."""
 
+    aws_hook_class = RdsHook
     ui_color = "#ddbb77"
     ui_fgcolor = "#ffffff"
 
-    def __init__(
-        self, *args, aws_conn_id: str | None = "aws_conn_id", hook_params: dict | None = None, **kwargs
-    ):
+    def __init__(self, *args, hook_params: dict | None = None, **kwargs):
         self.hook_params = hook_params or {}
-        self.aws_conn_id = aws_conn_id
         self.target_statuses: list[str] = []
         super().__init__(*args, **kwargs)
 
-    @cached_property
-    def hook(self):
-        return RdsHook(aws_conn_id=self.aws_conn_id, **self.hook_params)
-
 
 class RdsSnapshotExistenceSensor(RdsBaseSensor):
     """
@@ -59,9 +53,19 @@ class RdsSnapshotExistenceSensor(RdsBaseSensor):
     :param db_type: Type of the DB - either "instance" or "cluster"
     :param db_snapshot_identifier: The identifier for the DB snapshot
     :param target_statuses: Target status of snapshot
+    :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
     """
 
-    template_fields: Sequence[str] = (
+    template_fields: Sequence[str] = aws_template_fields(
         "db_snapshot_identifier",
         "target_statuses",
     )
@@ -72,10 +76,9 @@ class RdsSnapshotExistenceSensor(RdsBaseSensor):
         db_type: str,
         db_snapshot_identifier: str,
         target_statuses: list[str] | None = None,
-        aws_conn_id: str | None = "aws_conn_id",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
         self.db_type = RdsDbType(db_type)
         self.db_snapshot_identifier = db_snapshot_identifier
         self.target_statuses = target_statuses or ["available"]
@@ -107,7 +110,9 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
     :param error_statuses: Target error status of export task to fail the sensor
     """
 
-    template_fields: Sequence[str] = ("export_task_identifier", "target_statuses", "error_statuses")
+    template_fields: Sequence[str] = aws_template_fields(
+        "export_task_identifier", "target_statuses", "error_statuses"
+    )
 
     def __init__(
         self,
@@ -115,10 +120,9 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
         export_task_identifier: str,
         target_statuses: list[str] | None = None,
         error_statuses: list[str] | None = None,
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.export_task_identifier = export_task_identifier
         self.target_statuses = target_statuses or [
@@ -159,7 +163,7 @@ class RdsDbSensor(RdsBaseSensor):
     :param target_statuses: Target status of DB
     """
 
-    template_fields: Sequence[str] = (
+    template_fields: Sequence[str] = aws_template_fields(
         "db_identifier",
         "db_type",
         "target_statuses",
@@ -171,10 +175,9 @@ class RdsDbSensor(RdsBaseSensor):
         db_identifier: str,
         db_type: RdsDbType | str = RdsDbType.INSTANCE,
         target_statuses: list[str] | None = None,
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
         self.db_identifier = db_identifier
         self.target_statuses = target_statuses or ["available"]
         self.db_type = db_type

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

@@ -20,6 +20,7 @@ from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
@@ -182,3 +183,100 @@ class BedrockIngestionJobTrigger(AwsBaseWaiterTrigger):
 
     def hook(self) -> AwsGenericHook:
         return BedrockAgentHook(aws_conn_id=self.aws_conn_id)
+
+
+class BedrockBaseBatchInferenceTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when a batch inference job is complete.
+
+    :param job_arn: The Amazon Resource Name (ARN) of the batch inference job.
+
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *,
+        job_arn: str,
+        waiter_name: str | ArgNotSet = NOTSET,  # This must be defined in the child class.
+        waiter_delay: int = 120,
+        waiter_max_attempts: int = 75,
+        aws_conn_id: str | None = None,
+    ) -> None:
+        if waiter_name == NOTSET:
+            raise NotImplementedError("Triggers must provide a waiter name.")
+
+        super().__init__(
+            serialized_fields={"job_arn": job_arn},
+            waiter_name=str(waiter_name),  # Cast a string to a string to make mypy happy
+            waiter_args={"jobIdentifier": job_arn},
+            failure_message="Bedrock batch inference job failed.",
+            status_message="Status of Bedrock batch inference job is",
+            status_queries=["status"],
+            return_key="job_arn",
+            return_value=job_arn,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return BedrockHook(aws_conn_id=self.aws_conn_id)
+
+
+class BedrockBatchInferenceCompletedTrigger(BedrockBaseBatchInferenceTrigger):
+    """
+    Trigger when a batch inference job is complete.
+
+    :param job_arn: The Amazon Resource Name (ARN) of the batch inference job.
+
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *,
+        job_arn: str,
+        waiter_delay: int = 120,
+        waiter_max_attempts: int = 75,
+        aws_conn_id: str | None = None,
+    ) -> None:
+        super().__init__(
+            waiter_name="batch_inference_complete",
+            job_arn=job_arn,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+
+class BedrockBatchInferenceScheduledTrigger(BedrockBaseBatchInferenceTrigger):
+    """
+    Trigger when a batch inference job is scheduled.
+
+    :param job_arn: The Amazon Resource Name (ARN) of the batch inference job.
+
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *,
+        job_arn: str,
+        waiter_delay: int = 120,
+        waiter_max_attempts: int = 75,
+        aws_conn_id: str | None = None,
+    ) -> None:
+        super().__init__(
+            waiter_name="batch_inference_scheduled",
+            job_arn=job_arn,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )

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

@@ -23,7 +23,7 @@ import time
 from typing import TYPE_CHECKING, Any
 
 import jmespath
-from botocore.exceptions import WaiterError
+from botocore.exceptions import NoCredentialsError, WaiterError
 
 from airflow.exceptions import AirflowException
 
@@ -70,6 +70,10 @@ def wait(
             time.sleep(waiter_delay)
         try:
             waiter.wait(**args, WaiterConfig={"MaxAttempts": 1})
+
+        except NoCredentialsError as error:
+            log.info(str(error))
+
         except WaiterError as error:
             error_reason = str(error)
             last_response = error.last_response
@@ -131,6 +135,10 @@ async def async_wait(
             await asyncio.sleep(waiter_delay)
         try:
             await waiter.wait(**args, WaiterConfig={"MaxAttempts": 1})
+
+        except NoCredentialsError as error:
+            log.info(str(error))
+
         except WaiterError as error:
             error_reason = str(error)
             last_response = error.last_response

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

@@ -68,6 +68,140 @@
                     "state": "failure"
                 }
             ]
+        },
+        "batch_inference_complete": {
+            "delay": 120,
+            "maxAttempts": 75,
+            "operation": "GetModelInvocationJob",
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Completed",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Failed",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Stopped",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "PartiallyCompleted",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Expired",
+                    "state": "failure"
+                },
+                                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Stopping",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Submitted",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "InProgress",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Validating",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Scheduled",
+                    "state": "retry"
+                }
+            ]
+        },
+        "batch_inference_scheduled": {
+            "delay": 120,
+            "maxAttempts": 75,
+            "operation": "GetModelInvocationJob",
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Completed",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Failed",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Stopped",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Stopping",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "PartiallyCompleted",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Expired",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Submitted",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "InProgress",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Validating",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "status",
+                    "expected": "Scheduled",
+                    "state": "success"
+                }
+            ]
         }
     }
 }

airflow/providers/amazon/get_provider_info.py

@@ -26,79 +26,6 @@ def get_provider_info():
         "package-name": "apache-airflow-providers-amazon",
         "name": "Amazon",
         "description": "Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).\n",
-        "state": "ready",
-        "source-date-epoch": 1743477760,
-        "versions": [
-            "9.5.0",
-            "9.4.0",
-            "9.2.0",
-            "9.1.0",
-            "9.0.0",
-            "8.29.0",
-            "8.28.0",
-            "8.27.0",
-            "8.26.0",
-            "8.25.0",
-            "8.24.0",
-            "8.23.0",
-            "8.22.0",
-            "8.21.0",
-            "8.20.0",
-            "8.19.0",
-            "8.18.0",
-            "8.17.0",
-            "8.16.0",
-            "8.15.0",
-            "8.14.0",
-            "8.13.0",
-            "8.12.0",
-            "8.11.0",
-            "8.10.0",
-            "8.9.0",
-            "8.8.0",
-            "8.7.1",
-            "8.7.0",
-            "8.6.0",
-            "8.5.1",
-            "8.5.0",
-            "8.4.0",
-            "8.3.1",
-            "8.3.0",
-            "8.2.0",
-            "8.1.0",
-            "8.0.0",
-            "7.4.1",
-            "7.4.0",
-            "7.3.0",
-            "7.2.1",
-            "7.2.0",
-            "7.1.0",
-            "7.0.0",
-            "6.2.0",
-            "6.1.0",
-            "6.0.0",
-            "5.1.0",
-            "5.0.0",
-            "4.1.0",
-            "4.0.0",
-            "3.4.0",
-            "3.3.0",
-            "3.2.0",
-            "3.1.1",
-            "3.0.0",
-            "2.6.0",
-            "2.5.0",
-            "2.4.0",
-            "2.3.0",
-            "2.2.0",
-            "2.1.0",
-            "2.0.0",
-            "1.4.0",
-            "1.3.0",
-            "1.2.0",
-            "1.1.0",
-            "1.0.0",
-        ],
         "integrations": [
             {
                 "integration-name": "Amazon Athena",
@@ -1370,55 +1297,4 @@ def get_provider_info():
             },
         },
         "executors": ["airflow.providers.amazon.aws.executors.ecs.ecs_executor.AwsEcsExecutor"],
-        "dependencies": [
-            "apache-airflow>=2.9.0",
-            "apache-airflow-providers-common-compat>=1.6.0",
-            "apache-airflow-providers-common-sql>=1.20.0",
-            "apache-airflow-providers-http",
-            "boto3>=1.37.0",
-            "botocore>=1.34.90",
-            "inflection>=0.5.1",
-            "watchtower>=3.0.0,!=3.3.0,<4",
-            "jsonpath_ng>=1.5.3",
-            "redshift_connector>=2.0.918",
-            "asgiref>=2.3.0",
-            "PyAthena>=3.0.10",
-            "jmespath>=0.7.0",
-            "python3-saml>=1.16.0",
-            "xmlsec!=1.3.15,>=1.3.14",
-            "sagemaker-studio>=1.0.9",
-        ],
-        "optional-dependencies": {
-            "pandas": ["pandas>=2.1.2,<2.2"],
-            "aiobotocore": ["aiobotocore[boto3]>=2.13.0"],
-            "cncf.kubernetes": ["apache-airflow-providers-cncf-kubernetes>=7.2.0"],
-            "s3fs": ["s3fs>=2023.10.0"],
-            "python3-saml": ["python3-saml>=1.16.0"],
-            "apache.hive": ["apache-airflow-providers-apache-hive"],
-            "exasol": ["apache-airflow-providers-exasol"],
-            "fab": ["apache-airflow-providers-fab"],
-            "ftp": ["apache-airflow-providers-ftp"],
-            "google": ["apache-airflow-providers-google"],
-            "imap": ["apache-airflow-providers-imap"],
-            "microsoft.azure": ["apache-airflow-providers-microsoft-azure"],
-            "mongo": ["apache-airflow-providers-mongo"],
-            "openlineage": ["apache-airflow-providers-openlineage"],
-            "salesforce": ["apache-airflow-providers-salesforce"],
-            "ssh": ["apache-airflow-providers-ssh"],
-            "standard": ["apache-airflow-providers-standard"],
-        },
-        "devel-dependencies": [
-            "aiobotocore>=2.13.0",
-            "aws_xray_sdk>=2.12.0",
-            "moto[cloudformation,glue]>=5.1.2",
-            "mypy-boto3-appflow>=1.37.0",
-            "mypy-boto3-rds>=1.34.90",
-            "mypy-boto3-redshift-data>=1.34.0",
-            "mypy-boto3-s3>=1.34.90",
-            "s3fs>=2023.10.0",
-            "openapi-schema-validator>=0.6.2",
-            "openapi-spec-validator>=0.7.1",
-            "opensearch-py>=2.2.0",
-            "responses>=0.25.0",
-        ],
     }

{apache_airflow_providers_amazon-9.5.0rc3.dist-info → apache_airflow_providers_amazon-9.6.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-amazon
-Version: 9.5.0rc3
+Version: 9.6.0rc1
 Summary: Provider package apache-airflow-providers-amazon for Apache Airflow
 Keywords: airflow-provider,amazon,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -25,18 +25,18 @@ Requires-Dist: apache-airflow-providers-common-compat>=1.6.0rc0
 Requires-Dist: apache-airflow-providers-common-sql>=1.20.0rc0
 Requires-Dist: apache-airflow-providers-http
 Requires-Dist: boto3>=1.37.0
-Requires-Dist: botocore>=1.34.90
+Requires-Dist: botocore>=1.37.0
 Requires-Dist: inflection>=0.5.1
-Requires-Dist: watchtower>=3.0.0,!=3.3.0,<4
+Requires-Dist: watchtower>=3.3.1,<4
 Requires-Dist: jsonpath_ng>=1.5.3
-Requires-Dist: redshift_connector>=2.0.918
+Requires-Dist: redshift_connector>=2.1.3
 Requires-Dist: asgiref>=2.3.0
-Requires-Dist: PyAthena>=3.0.10
+Requires-Dist: PyAthena>=3.10.0
 Requires-Dist: jmespath>=0.7.0
 Requires-Dist: python3-saml>=1.16.0
 Requires-Dist: xmlsec!=1.3.15,>=1.3.14
 Requires-Dist: sagemaker-studio>=1.0.9
-Requires-Dist: aiobotocore[boto3]>=2.13.0 ; extra == "aiobotocore"
+Requires-Dist: aiobotocore[boto3]>=2.20.0 ; extra == "aiobotocore"
 Requires-Dist: apache-airflow-providers-apache-hive ; extra == "apache-hive"
 Requires-Dist: apache-airflow-providers-cncf-kubernetes>=7.2.0 ; extra == "cncf-kubernetes"
 Requires-Dist: apache-airflow-providers-exasol ; extra == "exasol"
@@ -54,8 +54,8 @@ Requires-Dist: apache-airflow-providers-salesforce ; extra == "salesforce"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Requires-Dist: apache-airflow-providers-standard ; extra == "standard"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.5.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.5.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.6.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.6.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -103,7 +103,7 @@ Provides-Extra: standard
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``9.5.0``
+Release: ``9.6.0``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -116,7 +116,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.5.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.6.0/>`_.
 
 Installation
 ------------
@@ -130,26 +130,26 @@ The package supports the following python versions: 3.9,3.10,3.11,3.12
 Requirements
 ------------
 
-==========================================  ======================
+==========================================  =====================
 PIP package                                 Version required
-==========================================  ======================
+==========================================  =====================
 ``apache-airflow``                          ``>=2.9.0``
 ``apache-airflow-providers-common-compat``  ``>=1.6.0``
 ``apache-airflow-providers-common-sql``     ``>=1.20.0``
 ``apache-airflow-providers-http``
 ``boto3``                                   ``>=1.37.0``
-``botocore``                                ``>=1.34.90``
+``botocore``                                ``>=1.37.0``
 ``inflection``                              ``>=0.5.1``
-``watchtower``                              ``>=3.0.0,!=3.3.0,<4``
+``watchtower``                              ``>=3.3.1,<4``
 ``jsonpath_ng``                             ``>=1.5.3``
-``redshift_connector``                      ``>=2.0.918``
+``redshift_connector``                      ``>=2.1.3``
 ``asgiref``                                 ``>=2.3.0``
-``PyAthena``                                ``>=3.0.10``
+``PyAthena``                                ``>=3.10.0``
 ``jmespath``                                ``>=0.7.0``
 ``python3-saml``                            ``>=1.16.0``
 ``xmlsec``                                  ``>=1.3.14,!=1.3.15``
 ``sagemaker-studio``                        ``>=1.0.9``
-==========================================  ======================
+==========================================  =====================
 
 Cross provider package dependencies
 -----------------------------------
@@ -184,5 +184,5 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.5.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.6.0/changelog.html>`_.