diff --git a/providers/amazon/docs/changelog.rst b/providers/amazon/docs/changelog.rst index 214f844c59263..fa3f4e9ee49fe 100644 --- a/providers/amazon/docs/changelog.rst +++ b/providers/amazon/docs/changelog.rst @@ -26,6 +26,15 @@ Changelog --------- +.. warning:: + Deferrable AWS operators and sensors now hand ``region_name``, ``verify`` and ``botocore_config`` + to the trigger they defer to, so the triggerer builds its hook from the operator's settings + instead of falling back to boto3 defaults. Deployments where the triggerer happened to work + *because* of those defaults will see it change: it now uses the operator's region rather than the + triggerer host's ``AWS_DEFAULT_REGION``, and it applies the operator's SSL verification and + botocore configuration, which previously never reached it. Set these explicitly on the operator + if the deferred half needs to differ from the synchronous half. + 9.36.0 ...... diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py index 62e49872f658f..3e0f0bc0b7aff 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py @@ -211,6 +211,9 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -393,6 +396,9 @@ def execute(self, context: Context) -> None: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -545,6 +551,9 @@ def execute(self, context: Context) -> dict: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -634,6 +643,9 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -807,6 +819,9 @@ def _create_kb(): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -994,6 +1009,9 @@ def start_ingestion_job(): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -1311,6 +1329,9 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/comprehend.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/comprehend.py index cdc3a30bf9e4c..6cdf3e175b513 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/comprehend.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/comprehend.py @@ -192,6 +192,9 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -341,6 +344,9 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py index ecf83773421c5..2c615cbc0794c 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py @@ -263,6 +263,9 @@ def execute(self, context: Context) -> dict: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", kwargs={"result": result}, @@ -765,6 +768,9 @@ def execute(self, context: Context) -> None: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="retry_execution", ) @@ -776,6 +782,9 @@ def execute(self, context: Context) -> None: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="retry_execution", ) @@ -797,6 +806,9 @@ def handle_delete_wait(self): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -951,6 +963,9 @@ def execute(self, context: Context): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="retry_execution", ) @@ -994,6 +1009,9 @@ def execute(self, context: Context): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -1101,6 +1119,9 @@ def execute(self, context: Context) -> None: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py index 89fc67b6d0411..be6c171a56a3a 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py @@ -141,6 +141,8 @@ def execute(self, context: Context): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, ), method_name="_complete_exec_with_cluster_desc", @@ -218,6 +220,8 @@ def execute(self, context: Context): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, ), method_name="_complete_exec_with_cluster_desc", @@ -622,7 +626,10 @@ def execute_complete(self, context: Context, event: dict[str, Any] | None = None if self._aws_logs_enabled(): # same behavior as non-deferrable mode, return last line of logs of the task. logs_client = AwsLogsHook( - aws_conn_id=self.aws_conn_id, region_name=self.resolve_awslogs_region() + aws_conn_id=self.aws_conn_id, + region_name=self.resolve_awslogs_region(), + verify=self.verify, + config=self.botocore_config, ).conn one_log = logs_client.get_log_events( logGroupName=self.awslogs_group, diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py index c2efc42f8a57b..e4fa2efde429d 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py @@ -359,6 +359,8 @@ def execute(self, context: Context): trigger=EksCreateClusterTrigger( cluster_name=self.cluster_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -411,6 +413,8 @@ def deferrable_create_cluster_next(self, context: Context, event: dict[str, Any] waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, force_delete_compute=False, ), @@ -445,6 +449,8 @@ def deferrable_create_cluster_next(self, context: Context, event: dict[str, Any] waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, ), method_name="execute_complete", @@ -456,6 +462,8 @@ def deferrable_create_cluster_next(self, context: Context, event: dict[str, Any] nodegroup_name=self.nodegroup_name, cluster_name=self.cluster_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -600,6 +608,8 @@ def execute(self, context: Context): cluster_name=self.cluster_name, nodegroup_name=self.nodegroup_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -714,6 +724,8 @@ def execute(self, context: Context): cluster_name=self.cluster_name, fargate_profile_name=self.fargate_profile_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, region_name=self.region_name, @@ -805,6 +817,8 @@ def execute(self, context: Context): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, force_delete_compute=self.force_delete_compute, ), @@ -946,6 +960,8 @@ def execute(self, context: Context): cluster_name=self.cluster_name, nodegroup_name=self.nodegroup_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -1038,6 +1054,8 @@ def execute(self, context: Context): cluster_name=self.cluster_name, fargate_profile_name=self.fargate_profile_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, region_name=self.region_name, diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py index 20366b694877c..2c66df54f59a2 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py @@ -236,6 +236,9 @@ def execute(self, context: Context) -> list[str]: job_flow_id=job_flow_id, step_ids=step_ids, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_max_attempts=self.waiter_max_attempts, waiter_delay=self.waiter_delay, ), @@ -645,6 +648,9 @@ def execute(self, context: Context) -> str | None: virtual_cluster_id=self.virtual_cluster_id, job_id=self.job_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.poll_interval, waiter_max_attempts=self.max_polling_attempts, cancel_on_kill=self.cancel_on_kill, @@ -654,6 +660,9 @@ def execute(self, context: Context) -> str | None: virtual_cluster_id=self.virtual_cluster_id, job_id=self.job_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.poll_interval, cancel_on_kill=self.cancel_on_kill, ), @@ -866,6 +875,9 @@ def execute(self, context: Context) -> str | None: trigger=EmrCreateJobFlowTrigger( job_flow_id=self._job_flow_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, waiter_name=waiter_name, @@ -1079,6 +1091,9 @@ def execute(self, context: Context) -> None: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, 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 is set to ensure that if a trigger dies, the timeout does not restart @@ -1175,6 +1190,9 @@ def execute(self, context: Context) -> str | None: trigger=EmrServerlessCreateApplicationTrigger( application_id=application_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, ), @@ -1220,6 +1238,9 @@ def start_application_deferred(self, context: Context, event: dict[str, Any] | N trigger=EmrServerlessStartApplicationTrigger( application_id=event["application_id"], aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, ), @@ -1368,6 +1389,9 @@ def execute(self, context: Context, event: dict[str, Any] | None = None) -> str waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute", timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay), @@ -1427,6 +1451,9 @@ def execute(self, context: Context, event: dict[str, Any] | None = None) -> str waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, cancel_on_kill=self.cancel_on_kill, ), method_name="execute_complete", @@ -1686,6 +1713,9 @@ def execute(self, context: Context) -> None: trigger=EmrServerlessCancelJobsTrigger( application_id=self.application_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, ), @@ -1709,6 +1739,9 @@ def execute(self, context: Context) -> None: trigger=EmrServerlessStopApplicationTrigger( application_id=self.application_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, ), @@ -1739,6 +1772,9 @@ def stop_application(self, context: Context, event: dict[str, Any] | None = None trigger=EmrServerlessStopApplicationTrigger( application_id=self.application_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, ), @@ -1842,6 +1878,9 @@ def _delete_application(self) -> None: trigger=EmrServerlessDeleteApplicationTrigger( application_id=self.application_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, ), diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py index 6631706a65446..ca2ff750d300e 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py @@ -756,6 +756,9 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -904,6 +907,9 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py index d9cd6cbdf38d6..ee468f7a4cdfd 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py @@ -170,6 +170,9 @@ def execute(self, context: Context) -> dict: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py index 134d13df982c5..f8cc41881bd48 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py @@ -82,6 +82,8 @@ class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]): 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 :return: dictionary with Neptune graph id @@ -173,6 +175,9 @@ def execute(self, context: Context) -> dict: self.defer( trigger=NeptuneGraphAvailableTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_id, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -230,6 +235,8 @@ class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalytics 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 :return: dictionary with Neptune graph id @@ -317,6 +324,9 @@ def execute(self, context: Context) -> dict: self.defer( trigger=NeptuneGraphPrivateEndpointAvailableTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_identifier, vpc_id=self.vpc_id, waiter_delay=self.waiter_delay, @@ -373,7 +383,8 @@ class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalytics 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 :return: dictionary with Neptune graph id @@ -420,6 +431,9 @@ def execute(self, context: Context) -> None: self.defer( trigger=NeptuneGraphPrivateEndpointDeletedTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_identifier, vpc_id=self.vpc_id, endpoint_id=endpoint_id, @@ -473,7 +487,8 @@ class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]): 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 :return: dictionary with Neptune graph id @@ -517,6 +532,9 @@ def execute(self, context: Context): self.defer( trigger=NeptuneGraphDeletedTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_id, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -582,6 +600,8 @@ class NeptuneCreateGraphWithImportOperator(AwsBaseOperator[NeptuneAnalyticsHook] 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 :return: dictionary with Neptune graph id @@ -729,6 +749,9 @@ def execute(self, context: Context) -> dict: self.defer( trigger=NeptuneGraphAvailableTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_id, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -773,6 +796,9 @@ def defer_wait_for_task( waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", kwargs={"graph_id": graph_id}, @@ -822,6 +848,8 @@ class NeptuneStartImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]): 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 :return: dictionary with Neptune graph id @@ -914,6 +942,9 @@ def execute(self, context: Context) -> dict: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -962,6 +993,8 @@ class NeptuneCancelImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]): 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 :return: dictionary with Neptune graph id @@ -1002,6 +1035,9 @@ def execute(self, context: Context) -> dict: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/rds.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/rds.py index 384547676a828..3a85c52edd2d4 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/rds.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/rds.py @@ -649,6 +649,8 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, # ignoring type because create_db_instance is a dict response=create_db_instance, # type: ignore[arg-type] @@ -739,6 +741,8 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, # ignoring type because delete_db_instance is a dict response=delete_db_instance, # type: ignore[arg-type] @@ -823,6 +827,8 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, response=start_db_response, db_type=self.db_type, @@ -927,6 +933,8 @@ def execute(self, context: Context) -> str: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, response=stop_db_response, db_type=self.db_type, diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py index e65a35dd62f42..d8c980e6ca34e 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py @@ -158,6 +158,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="poke", ) @@ -228,6 +231,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="poke", ) @@ -297,6 +303,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="poke", ) @@ -386,6 +395,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="poke", ) @@ -493,6 +505,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="poke", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/comprehend.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/comprehend.py index 98361710bb6de..6de2682d0ba48 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/comprehend.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/comprehend.py @@ -135,6 +135,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="poke", ) @@ -220,6 +223,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="poke", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/emr.py index 8e2536e78810b..dfc7c34af81c2 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/emr.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/emr.py @@ -357,6 +357,9 @@ def execute(self, context: Context): virtual_cluster_id=self.virtual_cluster_id, job_id=self.job_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.poll_interval, waiter_max_attempts=self.max_retries, ) @@ -365,6 +368,9 @@ def execute(self, context: Context): virtual_cluster_id=self.virtual_cluster_id, job_id=self.job_id, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.poll_interval, ), method_name="execute_complete", @@ -554,6 +560,9 @@ def execute(self, context: Context) -> None: job_flow_id=self.job_flow_id, waiter_max_attempts=self.max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=int(self.poke_interval), ), method_name="execute_complete", @@ -684,6 +693,9 @@ def execute(self, context: Context) -> None: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/glue.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/glue.py index a434a174cfef2..58190131c190f 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/glue.py @@ -212,6 +212,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -332,6 +335,9 @@ def execute(self, context: Context) -> Any: waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py index 37710a7f0bf0c..fb0fe4ea5df7d 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py @@ -161,6 +161,9 @@ def execute(self, context: Context): waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -312,6 +315,9 @@ def execute(self, context: Context): waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/ssm.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/ssm.py index 943c960873095..d4544c7b7369f 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/ssm.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/ssm.py @@ -142,6 +142,9 @@ def execute(self, context: Context): waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, fail_on_nonzero_exit=self.fail_on_nonzero_exit, ), method_name="execute_complete", diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py index b952e27fbbaf2..97f5b2c375014 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py @@ -17,24 +17,25 @@ from __future__ import annotations -from abc import abstractmethod from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any +from typing import Any from airflow.exceptions import AirflowException +from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.utils.helpers import prune_dict -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - class AwsBaseWaiterTrigger(BaseTrigger): """ Base class for all AWS Triggers that follow the "standard" model of just waiting on a waiter. - Subclasses need to implement the hook() method. + Subclasses should set the ``aws_hook_class`` attribute to the hook they need. The hook is then + built from the parameters this class already serializes, so the deferred half of a task talks to + AWS with the same region, SSL verification setting and botocore configuration as the synchronous + half. Subclasses whose hook takes something else may override :meth:`_hook_parameters` or, as a + last resort, :meth:`hook` itself. :param serialized_fields: Fields that are specific to the subclass trigger and need to be serialized to be passed to the __init__ method on deserialization. @@ -67,6 +68,34 @@ class AwsBaseWaiterTrigger(BaseTrigger): https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + # Should be assigned in child class, unless hook() is overridden. + aws_hook_class: type[AwsGenericHook] + + def __init_subclass__(cls, **kwargs: Any) -> None: + """ + Reject a subclass that cannot build a hook, at import time. + + ``hook()`` is only reached from ``run()``, which executes in the triggerer, so without this + a subclass that declares neither would defer successfully and fail later, out of sight of + the task that deferred. The operator side gets the same guarantee from + ``AwsBaseHookMixin.validate_attributes``. This runs on class creation rather than in + ``__init__`` because subclasses such as ``EksDeleteClusterTrigger`` never call + ``super().__init__()``. + """ + super().__init_subclass__(**kwargs) + if cls.hook is not AwsBaseWaiterTrigger.hook: + return + hook_class = getattr(cls, "aws_hook_class", None) + if hook_class is None: + raise AttributeError( + f"Class attribute '{cls.__name__}.aws_hook_class' should be set, " + f"or {cls.__name__}.hook() overridden." + ) + if not (isinstance(hook_class, type) and issubclass(hook_class, AwsGenericHook)): + raise AttributeError( + f"Class attribute '{cls.__name__}.aws_hook_class' is not a subclass of AwsGenericHook." + ) + def __init__( self, *, @@ -137,9 +166,24 @@ def serialize(self) -> tuple[str, dict[str, Any]]: params, ) - @abstractmethod + @property + def _hook_parameters(self) -> dict[str, Any]: + """Mapping of the serialized parameters onto the hook's constructor keywords.""" + return { + "aws_conn_id": self.aws_conn_id, + "region_name": self.region_name, + "verify": self.verify, + "config": self.botocore_config, + } + def hook(self) -> AwsGenericHook: - """Override in subclasses to return the right hook.""" + """Build the hook this trigger waits with.""" + if not hasattr(self, "aws_hook_class"): + raise AttributeError( + f"Class attribute '{type(self).__name__}.aws_hook_class' should be set, " + f"or {type(self).__name__}.hook() overridden." + ) + return self.aws_hook_class(**self._hook_parameters) def _event_from_exception(self, error: AirflowException) -> TriggerEvent: return TriggerEvent( diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py index 4671d62007eaa..20385aa815ff4 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py @@ -16,8 +16,6 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING - from airflow.providers.amazon.aws.hooks.bedrock import ( BedrockAgentCoreControlHook, BedrockAgentHook, @@ -26,9 +24,6 @@ from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - class BedrockCustomizeModelCompletedTrigger(AwsBaseWaiterTrigger): """ @@ -38,8 +33,15 @@ class BedrockCustomizeModelCompletedTrigger(AwsBaseWaiterTrigger): :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. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = BedrockHook + def __init__( self, *, @@ -47,6 +49,9 @@ def __init__( waiter_delay: int = 120, waiter_max_attempts: int = 75, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"job_name": job_name}, @@ -60,11 +65,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return BedrockHook(aws_conn_id=self.aws_conn_id) - class BedrockKnowledgeBaseActiveTrigger(AwsBaseWaiterTrigger): """ @@ -75,8 +80,15 @@ class BedrockKnowledgeBaseActiveTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 5) :param waiter_max_attempts: The maximum number of attempts to be made. (default: 24) :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = BedrockAgentHook + def __init__( self, *, @@ -84,6 +96,9 @@ def __init__( waiter_delay: int = 5, waiter_max_attempts: int = 24, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"knowledge_base_id": knowledge_base_id}, @@ -97,11 +112,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return BedrockAgentHook(aws_conn_id=self.aws_conn_id) - class BedrockProvisionModelThroughputCompletedTrigger(AwsBaseWaiterTrigger): """ @@ -112,8 +127,15 @@ class BedrockProvisionModelThroughputCompletedTrigger(AwsBaseWaiterTrigger): :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. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = BedrockHook + def __init__( self, *, @@ -121,6 +143,9 @@ def __init__( waiter_delay: int = 120, waiter_max_attempts: int = 75, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"provisioned_model_id": provisioned_model_id}, @@ -134,11 +159,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return BedrockHook(aws_conn_id=self.aws_conn_id) - class BedrockIngestionJobTrigger(AwsBaseWaiterTrigger): """ @@ -151,8 +176,15 @@ class BedrockIngestionJobTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60) :param waiter_max_attempts: The maximum number of attempts to be made. (default: 10) :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = BedrockAgentHook + def __init__( self, *, @@ -162,6 +194,9 @@ def __init__( waiter_delay: int = 60, waiter_max_attempts: int = 10, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={ @@ -183,11 +218,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return BedrockAgentHook(aws_conn_id=self.aws_conn_id) - class BedrockAgentRuntimeReadyTrigger(AwsBaseWaiterTrigger): """ @@ -199,8 +234,15 @@ class BedrockAgentRuntimeReadyTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60) :param waiter_max_attempts: The maximum number of attempts to be made. (default: 20) :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = BedrockAgentCoreControlHook + def __init__( self, *, @@ -210,6 +252,9 @@ def __init__( waiter_delay: int = 60, waiter_max_attempts: int = 20, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={ @@ -230,11 +275,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return BedrockAgentCoreControlHook(aws_conn_id=self.aws_conn_id) - class BedrockAgentRuntimeDeletedTrigger(AwsBaseWaiterTrigger): """ @@ -244,8 +289,15 @@ class BedrockAgentRuntimeDeletedTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60) :param waiter_max_attempts: The maximum number of attempts to be made. (default: 20) :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = BedrockAgentCoreControlHook + def __init__( self, *, @@ -253,6 +305,9 @@ def __init__( waiter_delay: int = 60, waiter_max_attempts: int = 20, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"agent_runtime_id": agent_runtime_id}, @@ -266,11 +321,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return BedrockAgentCoreControlHook(aws_conn_id=self.aws_conn_id) - class BedrockBaseBatchInferenceTrigger(AwsBaseWaiterTrigger): """ @@ -281,8 +336,15 @@ class BedrockBaseBatchInferenceTrigger(AwsBaseWaiterTrigger): :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. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = BedrockHook + def __init__( self, *, @@ -291,6 +353,9 @@ def __init__( waiter_delay: int = 120, waiter_max_attempts: int = 75, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: if waiter_name == NOTSET: raise NotImplementedError("Triggers must provide a waiter name.") @@ -307,11 +372,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return BedrockHook(aws_conn_id=self.aws_conn_id) - class BedrockBatchInferenceCompletedTrigger(BedrockBaseBatchInferenceTrigger): """ @@ -322,6 +387,11 @@ class BedrockBatchInferenceCompletedTrigger(BedrockBaseBatchInferenceTrigger): :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. + :param region_name: The AWS region where the resources to watch are. + :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__( @@ -331,6 +401,9 @@ def __init__( waiter_delay: int = 120, waiter_max_attempts: int = 75, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( waiter_name="batch_inference_complete", @@ -338,6 +411,9 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) @@ -350,6 +426,11 @@ class BedrockBatchInferenceScheduledTrigger(BedrockBaseBatchInferenceTrigger): :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. + :param region_name: The AWS region where the resources to watch are. + :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__( @@ -359,6 +440,9 @@ def __init__( waiter_delay: int = 120, waiter_max_attempts: int = 75, aws_conn_id: str | None = None, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( waiter_name="batch_inference_scheduled", @@ -366,4 +450,7 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/comprehend.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/comprehend.py index b2606813e69c6..5fcd64ed19a36 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/comprehend.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/comprehend.py @@ -16,11 +16,6 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - from airflow.providers.amazon.aws.hooks.comprehend import ComprehendHook from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger @@ -33,8 +28,15 @@ class ComprehendPiiEntitiesDetectionJobCompletedTrigger(AwsBaseWaiterTrigger): :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. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = ComprehendHook + def __init__( self, *, @@ -42,6 +44,9 @@ def __init__( waiter_delay: int = 120, waiter_max_attempts: int = 75, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"job_id": job_id}, @@ -55,11 +60,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return ComprehendHook(aws_conn_id=self.aws_conn_id) - class ComprehendCreateDocumentClassifierCompletedTrigger(AwsBaseWaiterTrigger): """ @@ -69,8 +74,15 @@ class ComprehendCreateDocumentClassifierCompletedTrigger(AwsBaseWaiterTrigger): :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. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = ComprehendHook + def __init__( self, *, @@ -78,6 +90,9 @@ def __init__( waiter_delay: int = 120, waiter_max_attempts: int = 75, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"document_classifier_arn": document_classifier_arn}, @@ -91,7 +106,7 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - - def hook(self) -> AwsGenericHook: - return ComprehendHook(aws_conn_id=self.aws_conn_id) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/dms.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/dms.py index 9f784de3c4092..921f551bc307e 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/dms.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/dms.py @@ -17,18 +17,14 @@ from __future__ import annotations from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any +from typing import Any from airflow.exceptions import AirflowException -from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook from airflow.providers.amazon.aws.hooks.dms import DmsHook from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait from airflow.triggers.base import BaseTrigger, TriggerEvent -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - class DmsReplicationTerminalStatusTrigger(AwsBaseWaiterTrigger): """ @@ -38,14 +34,24 @@ class DmsReplicationTerminalStatusTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = DmsHook + def __init__( self, replication_config_arn: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"replication_config_arn": replication_config_arn}, @@ -59,13 +65,9 @@ def __init__( return_key="replication_config_arn", return_value=replication_config_arn, aws_conn_id=aws_conn_id, - ) - - def hook(self) -> AwsGenericHook: - return DmsHook( - self.aws_conn_id, - verify=self.verify, - config=self.botocore_config, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) @@ -77,14 +79,24 @@ class DmsReplicationConfigDeletedTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = DmsHook + def __init__( self, replication_config_arn: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"replication_config_arn": replication_config_arn}, @@ -98,13 +110,9 @@ def __init__( return_key="replication_config_arn", return_value=replication_config_arn, aws_conn_id=aws_conn_id, - ) - - def hook(self) -> AwsGenericHook: - return DmsHook( - self.aws_conn_id, - verify=self.verify, - config=self.botocore_config, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) @@ -116,14 +124,24 @@ class DmsReplicationCompleteTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = DmsHook + def __init__( self, replication_config_arn: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"replication_config_arn": replication_config_arn}, @@ -137,13 +155,9 @@ def __init__( return_key="replication_config_arn", return_value=replication_config_arn, aws_conn_id=aws_conn_id, - ) - - def hook(self) -> AwsGenericHook: - return DmsHook( - self.aws_conn_id, - verify=self.verify, - config=self.botocore_config, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) @@ -155,14 +169,24 @@ class DmsReplicationStoppedTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = DmsHook + def __init__( self, replication_config_arn: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"replication_config_arn": replication_config_arn}, @@ -176,13 +200,9 @@ def __init__( return_key="replication_config_arn", return_value=replication_config_arn, aws_conn_id=aws_conn_id, - ) - - def hook(self) -> AwsGenericHook: - return DmsHook( - self.aws_conn_id, - verify=self.verify, - config=self.botocore_config, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) @@ -194,14 +214,24 @@ class DmsReplicationDeprovisionedTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = DmsHook + def __init__( self, replication_config_arn: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"replication_config_arn": replication_config_arn}, @@ -215,13 +245,9 @@ def __init__( return_key="replication_config_arn", return_value=replication_config_arn, aws_conn_id=aws_conn_id, - ) - - def hook(self) -> AwsGenericHook: - return DmsHook( - self.aws_conn_id, - verify=self.verify, - config=self.botocore_config, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) @@ -235,8 +261,11 @@ class DmsTaskModifyCompleteTrigger(AwsBaseWaiterTrigger): :param aws_conn_id: The Airflow connection used for AWS credentials. :param verify: Whether or not to verify SSL certificates. :param botocore_config: Configuration dictionary (key-values) for botocore client. + :param region_name: The AWS region where the resources to watch are. """ + aws_hook_class = DmsHook + def __init__( self, replication_task_arn: str, @@ -245,6 +274,7 @@ def __init__( aws_conn_id: str | None = "aws_default", verify: bool | str | None = None, botocore_config: dict | None = None, + region_name: str | None = None, ) -> None: super().__init__( serialized_fields={"replication_task_arn": replication_task_arn}, @@ -263,13 +293,7 @@ def __init__( aws_conn_id=aws_conn_id, verify=verify, botocore_config=botocore_config, - ) - - def hook(self) -> AwsGenericHook: - return DmsHook( - self.aws_conn_id, - verify=self.verify, - config=self.botocore_config, + region_name=region_name, ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/ecs.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/ecs.py index 630a2c1de55c1..4918c9a87a10e 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/ecs.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/ecs.py @@ -20,7 +20,7 @@ import asyncio import warnings from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any +from typing import Any from botocore.exceptions import ClientError, WaiterError @@ -32,9 +32,6 @@ from airflow.providers.common.compat.sdk import AirflowException from airflow.triggers.base import BaseTrigger, TriggerEvent -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - class ClusterActiveTrigger(AwsBaseWaiterTrigger): """ @@ -46,8 +43,14 @@ class ClusterActiveTrigger(AwsBaseWaiterTrigger): Will fail after that many unsuccessful attempts. :param aws_conn_id: The Airflow connection used for AWS credentials. :param region_name: The AWS region where the cluster is located. + :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 """ + aws_hook_class = EcsHook + def __init__( self, cluster_arn: str, @@ -55,6 +58,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, **kwargs, ): super().__init__( @@ -70,12 +75,11 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, **kwargs, ) - def hook(self) -> AwsGenericHook: - return EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class ClusterInactiveTrigger(AwsBaseWaiterTrigger): """ @@ -87,8 +91,14 @@ class ClusterInactiveTrigger(AwsBaseWaiterTrigger): Will fail after that many unsuccessful attempts. :param aws_conn_id: The Airflow connection used for AWS credentials. :param region_name: The AWS region where the cluster is located. + :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 """ + aws_hook_class = EcsHook + def __init__( self, cluster_arn: str, @@ -96,6 +106,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, **kwargs, ): super().__init__( @@ -110,12 +122,11 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, **kwargs, ) - def hook(self) -> AwsGenericHook: - return EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class TaskDoneTrigger(BaseTrigger): """ diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py index 18535d2344af2..1d142d9dbb645 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py @@ -29,12 +29,11 @@ from airflow.providers.cncf.kubernetes.triggers.pod import KubernetesPodTrigger from airflow.providers.common.compat.sdk import AirflowException from airflow.triggers.base import TriggerEvent +from airflow.utils.helpers import prune_dict if TYPE_CHECKING: from pendulum import DateTime - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - class EksCreateClusterTrigger(AwsBaseWaiterTrigger): """ @@ -48,8 +47,14 @@ class EksCreateClusterTrigger(AwsBaseWaiterTrigger): :param aws_conn_id: The Airflow connection used for AWS credentials. :param region_name: Which AWS region the connection should use. If this is None or empty 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 """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -57,6 +62,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"cluster_name": cluster_name, "region_name": region_name}, @@ -70,11 +77,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - async def run(self): async with await self.hook().get_async_conn() as client: waiter = client.get_waiter(self.waiter_name) @@ -237,8 +243,14 @@ class EksDeleteClusterTrigger(AwsBaseWaiterTrigger): If this is None or empty then the default boto3 behaviour is used. :param force_delete_compute: If True, any nodegroups or fargate profiles associated with the cluster will be deleted before the cluster is deleted. + :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 """ + aws_hook_class = EksHook + def __init__( self, cluster_name, @@ -247,12 +259,16 @@ def __init__( aws_conn_id: str | None, region_name: str | None, force_delete_compute: bool, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): self.cluster_name = cluster_name self.waiter_delay = waiter_delay self.waiter_max_attempts = waiter_max_attempts self.aws_conn_id = aws_conn_id self.region_name = region_name + self.verify = verify + self.botocore_config = botocore_config self.force_delete_compute = force_delete_compute def serialize(self) -> tuple[str, dict[str, Any]]: @@ -265,12 +281,10 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "aws_conn_id": self.aws_conn_id, "region_name": self.region_name, "force_delete_compute": self.force_delete_compute, + **prune_dict({"verify": self.verify, "botocore_config": self.botocore_config}), }, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - async def run(self): async with await self.hook().get_async_conn() as client: waiter = client.get_waiter("cluster_deleted") @@ -367,8 +381,15 @@ class EksCreateFargateProfileTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: Which AWS region the connection should use. + :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 """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -377,6 +398,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"cluster_name": cluster_name, "fargate_profile_name": fargate_profile_name}, @@ -390,11 +413,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class EksDeleteFargateProfileTrigger(AwsBaseWaiterTrigger): """ @@ -405,8 +427,15 @@ class EksDeleteFargateProfileTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: Which AWS region the connection should use. + :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 """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -415,6 +444,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"cluster_name": cluster_name, "fargate_profile_name": fargate_profile_name}, @@ -428,11 +459,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger): """ @@ -448,8 +478,14 @@ class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger): :param aws_conn_id: The Airflow connection used for AWS credentials. :param region_name: Which AWS region the connection should use. (templated) If this is None or empty 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 """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -458,6 +494,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={ @@ -475,11 +513,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class EksDeleteNodegroupTrigger(AwsBaseWaiterTrigger): """ @@ -495,8 +532,14 @@ class EksDeleteNodegroupTrigger(AwsBaseWaiterTrigger): :param aws_conn_id: The Airflow connection used for AWS credentials. :param region_name: Which AWS region the connection should use. (templated) If this is None or empty 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 """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -505,6 +548,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"cluster_name": cluster_name, "nodegroup_name": nodegroup_name}, @@ -518,7 +563,6 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py index bd545affd5709..8c557f494b1ff 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py @@ -19,7 +19,7 @@ import asyncio import sys from collections.abc import AsyncIterator -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from asgiref.sync import sync_to_async @@ -49,9 +49,16 @@ class EmrAddStepsTrigger(AwsBaseWaiterTrigger): :param waiter_delay: polling period in seconds to check for the status :param waiter_max_attempts: The maximum number of attempts to be made :param aws_conn_id: Reference to AWS connection id + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrHook + def __init__( self, job_flow_id: str, @@ -59,6 +66,9 @@ def __init__( waiter_delay: int, waiter_max_attempts: int, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"job_flow_id": job_flow_id, "step_ids": step_ids}, @@ -74,11 +84,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrHook(aws_conn_id=self.aws_conn_id) - class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger): """ @@ -88,8 +98,15 @@ class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrHook + def __init__( self, job_flow_id: str, @@ -97,6 +114,9 @@ def __init__( waiter_delay: int = 30, waiter_max_attempts: int = 60, waiter_name: str = "job_flow_waiting", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"job_flow_id": job_flow_id}, @@ -114,11 +134,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrHook(aws_conn_id=self.aws_conn_id) - class EmrTerminateJobFlowTrigger(AwsBaseWaiterTrigger): """ @@ -128,14 +148,24 @@ class EmrTerminateJobFlowTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrHook + def __init__( self, job_flow_id: str, aws_conn_id: str | None = None, waiter_delay: int = 30, waiter_max_attempts: int = 60, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"job_flow_id": job_flow_id}, @@ -152,11 +182,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrHook(aws_conn_id=self.aws_conn_id) - class EmrContainerTrigger(AwsBaseWaiterTrigger): """ @@ -171,8 +201,15 @@ class EmrContainerTrigger(AwsBaseWaiterTrigger): marks the deferred task failed, clears it, or mark-succeeds it. Requires ``apache-airflow`` with ``BaseTrigger.on_kill()`` support; on older versions the hook is silently inert. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrContainerHook + def __init__( self, virtual_cluster_id: str, @@ -181,6 +218,9 @@ def __init__( waiter_delay: int = 30, waiter_max_attempts: int = sys.maxsize, cancel_on_kill: bool = True, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={ @@ -198,13 +238,17 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) self.virtual_cluster_id = virtual_cluster_id self.job_id = job_id self.cancel_on_kill = cancel_on_kill - def hook(self) -> AwsGenericHook: - return EmrContainerHook(aws_conn_id=self.aws_conn_id, virtual_cluster_id=self.virtual_cluster_id) + @property + def _hook_parameters(self) -> dict[str, Any]: + return {**super()._hook_parameters, "virtual_cluster_id": self.virtual_cluster_id} async def on_kill(self) -> None: """Cancel the EMR container job when the user acts on the deferred task.""" @@ -235,8 +279,15 @@ class EmrStepSensorTrigger(AwsBaseWaiterTrigger): :param waiter_delay: polling period in seconds to check for the status :param waiter_max_attempts: The maximum number of attempts to be made :param aws_conn_id: Reference to AWS connection id + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrHook + def __init__( self, job_flow_id: str, @@ -244,6 +295,9 @@ def __init__( waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"job_flow_id": job_flow_id, "step_id": step_id}, @@ -260,11 +314,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrHook(aws_conn_id=self.aws_conn_id) - class EmrServerlessCreateApplicationTrigger(AwsBaseWaiterTrigger): """ @@ -274,14 +328,24 @@ class EmrServerlessCreateApplicationTrigger(AwsBaseWaiterTrigger): :waiter_delay: polling period in seconds to check for the status :param waiter_max_attempts: The maximum number of attempts to be made :param aws_conn_id: Reference to AWS connection id + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrServerlessHook + def __init__( self, application_id: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"application_id": application_id}, @@ -295,11 +359,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrServerlessHook(self.aws_conn_id) - class EmrServerlessStartApplicationTrigger(AwsBaseWaiterTrigger): """ @@ -309,14 +373,24 @@ class EmrServerlessStartApplicationTrigger(AwsBaseWaiterTrigger): :waiter_delay: polling period in seconds to check for the status :param waiter_max_attempts: The maximum number of attempts to be made :param aws_conn_id: Reference to AWS connection id + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrServerlessHook + def __init__( self, application_id: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"application_id": application_id}, @@ -330,11 +404,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrServerlessHook(self.aws_conn_id) - class EmrServerlessStopApplicationTrigger(AwsBaseWaiterTrigger): """ @@ -344,14 +418,24 @@ class EmrServerlessStopApplicationTrigger(AwsBaseWaiterTrigger): :waiter_delay: polling period in seconds to check for the status :param waiter_max_attempts: The maximum number of attempts to be made :param aws_conn_id: Reference to AWS connection id. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrServerlessHook + def __init__( self, application_id: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"application_id": application_id}, @@ -365,11 +449,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrServerlessHook(self.aws_conn_id) - class EmrServerlessJobSensorTrigger(AwsBaseWaiterTrigger): """ @@ -457,8 +541,15 @@ class EmrServerlessStartJobTrigger(AwsBaseWaiterTrigger): :param waiter_max_attempts: The maximum number of attempts to be made :param aws_conn_id: Reference to AWS connection id :param cancel_on_kill: Flag to indicate whether to cancel the job when the task is killed. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrServerlessHook + def __init__( self, application_id: str, @@ -467,6 +558,9 @@ def __init__( waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", cancel_on_kill: bool = True, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={ @@ -484,14 +578,14 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) self.application_id = application_id self.job_id = job_id self.cancel_on_kill = cancel_on_kill - def hook(self) -> AwsGenericHook: - return EmrServerlessHook(self.aws_conn_id) - if not AIRFLOW_V_3_0_PLUS: @provide_session @@ -633,14 +727,24 @@ class EmrServerlessDeleteApplicationTrigger(AwsBaseWaiterTrigger): :waiter_delay: polling period in seconds to check for the status :param waiter_max_attempts: The maximum number of attempts to be made :param aws_conn_id: Reference to AWS connection id + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrServerlessHook + def __init__( self, application_id: str, waiter_delay: int = 30, waiter_max_attempts: int = 60, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: super().__init__( serialized_fields={"application_id": application_id}, @@ -654,11 +758,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrServerlessHook(self.aws_conn_id) - class EmrServerlessCancelJobsTrigger(AwsBaseWaiterTrigger): """ @@ -668,14 +772,24 @@ class EmrServerlessCancelJobsTrigger(AwsBaseWaiterTrigger): :param aws_conn_id: Reference to AWS connection id :param waiter_delay: Delay in seconds between each attempt to check the status :param waiter_max_attempts: Maximum number of attempts to check the status + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = EmrServerlessHook + def __init__( self, application_id: str, aws_conn_id: str | None, waiter_delay: int, waiter_max_attempts: int, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: states = list(EmrServerlessHook.JOB_INTERMEDIATE_STATES.union({"CANCELLING"})) super().__init__( @@ -690,11 +804,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EmrServerlessHook(self.aws_conn_id) - @property def hook_instance(self) -> AwsGenericHook: """This property is added for backward compatibility.""" diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py index 04031294660a2..42c0e501decd8 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py @@ -20,13 +20,10 @@ import asyncio from collections.abc import AsyncIterator from functools import cached_property -from typing import TYPE_CHECKING, Any +from typing import Any from botocore.exceptions import ClientError -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - from airflow.providers.amazon.aws.hooks.glue import ( GlueDataQualityHook, GlueJobHook, @@ -55,6 +52,8 @@ class GlueJobCompleteTrigger(AwsBaseWaiterTrigger): :param botocore_config: Configuration dictionary (key-values) for botocore client. """ + aws_hook_class = GlueJobHook + def __init__( self, job_name: str, @@ -87,14 +86,6 @@ def __init__( self.run_id = run_id self.verbose = verbose - def hook(self) -> AwsGenericHook: - return GlueJobHook( - aws_conn_id=self.aws_conn_id, - region_name=self.region_name, - verify=self.verify, - config=self.botocore_config, - ) - async def run(self) -> AsyncIterator[TriggerEvent]: if not self.verbose: async for event in super().run(): @@ -105,7 +96,10 @@ async def run(self) -> AsyncIterator[TriggerEvent]: async with ( await hook.get_async_conn() as glue_client, await AwsLogsHook( - aws_conn_id=self.aws_conn_id, region_name=self.region_name + aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + config=self.botocore_config, ).get_async_conn() as logs_client, ): # Get log group names from job run metadata @@ -322,14 +316,24 @@ class GlueDataQualityRuleSetEvaluationRunCompleteTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60) :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. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = GlueDataQualityHook + def __init__( self, evaluation_run_id: str, waiter_delay: int = 60, waiter_max_attempts: int = 75, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"evaluation_run_id": evaluation_run_id}, @@ -343,11 +347,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return GlueDataQualityHook(aws_conn_id=self.aws_conn_id) - class GlueDataQualityRuleRecommendationRunCompleteTrigger(AwsBaseWaiterTrigger): """ @@ -357,14 +361,24 @@ class GlueDataQualityRuleRecommendationRunCompleteTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60) :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. + :param region_name: The AWS region where the resources to watch are. + :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 """ + aws_hook_class = GlueDataQualityHook + def __init__( self, recommendation_run_id: str, waiter_delay: int = 60, waiter_max_attempts: int = 75, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"recommendation_run_id": recommendation_run_id}, @@ -378,7 +392,7 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - - def hook(self) -> AwsGenericHook: - return GlueDataQualityHook(aws_conn_id=self.aws_conn_id) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/rds.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/rds.py index dd60037c69787..0e1ed8439bd6c 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/rds.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/rds.py @@ -16,16 +16,12 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any from airflow.providers.amazon.aws.hooks.rds import RdsHook from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger from airflow.providers.amazon.aws.utils.rds import RdsDbType -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - - _waiter_arg = { RdsDbType.INSTANCE.value: "DBInstanceIdentifier", RdsDbType.CLUSTER.value: "DBClusterIdentifier", @@ -47,8 +43,14 @@ class RdsDbAvailableTrigger(AwsBaseWaiterTrigger): :param region_name: AWS region where the DB is located, if different from the default one. :param response: The response from the RdsHook, to be passed back to the operator. :param db_type: The type of DB: instance or cluster. + :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 """ + aws_hook_class = RdsHook + def __init__( self, db_identifier: str, @@ -58,6 +60,8 @@ def __init__( response: dict[str, Any], db_type: RdsDbType | str, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: # allow passing enums for users, # but we can only rely on strings because (de-)serialization doesn't support enums @@ -83,11 +87,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class RdsDbDeletedTrigger(AwsBaseWaiterTrigger): """ @@ -100,8 +103,14 @@ class RdsDbDeletedTrigger(AwsBaseWaiterTrigger): :param region_name: AWS region where the DB is located, if different from the default one. :param response: The response from the RdsHook, to be passed back to the operator. :param db_type: The type of DB: instance or cluster. + :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 """ + aws_hook_class = RdsHook + def __init__( self, db_identifier: str, @@ -111,6 +120,8 @@ def __init__( response: dict[str, Any], db_type: RdsDbType | str, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: # allow passing enums for users, # but we can only rely on strings because (de-)serialization doesn't support enums @@ -136,11 +147,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class RdsDbStoppedTrigger(AwsBaseWaiterTrigger): """ @@ -153,8 +163,14 @@ class RdsDbStoppedTrigger(AwsBaseWaiterTrigger): :param region_name: AWS region where the DB is located, if different from the default one. :param response: The response from the RdsHook, to be passed back to the operator. :param db_type: The type of DB: instance or cluster. + :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 """ + aws_hook_class = RdsHook + def __init__( self, db_identifier: str, @@ -164,6 +180,8 @@ def __init__( response: dict[str, Any], db_type: RdsDbType | str, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: # allow passing enums for users, # but we can only rely on strings because (de-)serialization doesn't support enums @@ -189,7 +207,6 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - - def hook(self) -> AwsGenericHook: - return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py b/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py index 8a28069fd255e..137ab9c64c1e8 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py @@ -901,7 +901,42 @@ def test_execute_complete_uses_awslogs_region(self, check_mock, logs_hook_mock): assert result == "Log output" check_mock.assert_called_once_with() - logs_hook_mock.assert_called_once_with(aws_conn_id=self.ecs.aws_conn_id, region_name="logs-region") + logs_hook_mock.assert_called_once_with( + aws_conn_id=self.ecs.aws_conn_id, + region_name="logs-region", + verify=self.ecs.verify, + config=self.ecs.botocore_config, + ) + + @mock.patch("airflow.providers.amazon.aws.operators.ecs.AwsLogsHook") + @mock.patch.object(EcsRunTaskOperator, "_check_success_task") + def test_execute_complete_log_hook_uses_operator_aws_configuration(self, check_mock, logs_hook_mock): + botocore_config = {"read_timeout": 10} + self.set_up_operator( + awslogs_group="awslogs-group", + awslogs_region="logs-region", + awslogs_stream_prefix="prefix", + region_name="task-region", + verify="/path/to/ca-bundle.pem", + botocore_config=botocore_config, + ) + logs_hook_mock.return_value.conn.get_log_events.return_value = {"events": [{"message": "Log output"}]} + + self.ecs.execute_complete( + {}, + { + "status": "success", + "task_arn": f"arn:aws:ecs:us-east-1:012345678910:task/{TASK_ID}", + "cluster": "test_cluster", + }, + ) + + logs_hook_mock.assert_called_once_with( + aws_conn_id=self.ecs.aws_conn_id, + region_name="logs-region", + verify="/path/to/ca-bundle.pem", + config=botocore_config, + ) @mock.patch.object(EcsBaseOperator, "client") @mock.patch("airflow.providers.amazon.aws.utils.task_log_fetcher.AwsTaskLogFetcher") diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py index 566b4bee61384..4e91ab1e1b472 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py @@ -23,6 +23,7 @@ from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook from airflow.providers.amazon.aws.operators.mwaa import MwaaTriggerDagRunOperator +from airflow.providers.common.compat.sdk import TaskDeferred from unit.amazon.aws.utils.test_template_fields import validate_template_fields @@ -41,6 +42,9 @@ "waiter_max_attempts": 20, "deferrable": False, } +REGION_NAME = "eu-west-2" +VERIFY = False +BOTOCORE_CONFIG = {"read_timeout": 42} HOOK_RETURN_VALUE = { "ResponseMetadata": {}, "RestApiStatusCode": 200, @@ -115,3 +119,30 @@ def test_execute_wait_combinations(self, mock_hook, _, wait_for_completion, defe assert response == HOOK_RETURN_VALUE assert mock_hook.get_waiter.call_count == wait_for_completion assert op.defer.call_count == deferrable + + @mock.patch.object(MwaaTriggerDagRunOperator, "hook") + def test_deferred_trigger_receives_hook_configuration(self, mock_hook): + mock_hook.invoke_rest_api.return_value = HOOK_RETURN_VALUE + op = MwaaTriggerDagRunOperator( + **{**OP_KWARGS, "wait_for_completion": False, "deferrable": True}, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + op.execute({}) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": OP_KWARGS["waiter_delay"], + "waiter_max_attempts": OP_KWARGS["waiter_max_attempts"], + "aws_conn_id": "aws_default", + "external_env_name": OP_KWARGS["env_name"], + "external_dag_id": OP_KWARGS["trigger_dag_id"], + "external_dag_run_id": HOOK_RETURN_VALUE["RestApiResponse"]["dag_run_id"], + "success_states": None, + "failure_states": None, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py b/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py index 472e386888c31..ba92749660882 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py @@ -52,6 +52,9 @@ ENDPOINT_ID = "vpce-12345" SOURCE_S3_URI = "s3://my-bucket/my-data/" ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole" +REGION_NAME = "eu-west-2" +VERIFY = False +BOTOCORE_CONFIG = {"read_timeout": 42} class TestNeptuneCreateGraphOperator: @@ -219,6 +222,35 @@ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_pe assert isinstance(trigger, NeptuneGraphAvailableTrigger) assert exc_info.value.method_name == "execute_complete" + @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist") + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_persist): + mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"} + + operator = NeptuneCreateGraphOperator( + task_id="test_task", + graph_name=GRAPH_NAME, + vector_search_config={"test": 123}, + provisioned_memory=16, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } + class TestNeptuneCreatePrivateGraphEndpointOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -417,6 +449,40 @@ def test_execute_complete(self, mock_get_endpoint, mock_conn): ) assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID} + @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.VpcEndpointLink.persist") + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_persist): + mock_conn.create_private_graph_endpoint.return_value = { + "status": "CREATING", + "vpcEndpointId": ENDPOINT_ID, + "vpcId": VPC_ID, + } + mock_conn.get_private_graph_endpoint.return_value = {"vpcEndpointId": ENDPOINT_ID} + + operator = NeptuneCreatePrivateGraphEndpointOperator( + task_id="test_task", + graph_identifier=GRAPH_ID, + vpc_id=VPC_ID, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "vpc_id": VPC_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } + class TestNeptuneDeletePrivateGraphEndpointOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -548,6 +614,39 @@ def test_execute_complete_success(self): # Verify the method completes without error and logs the endpoint_id + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn): + mock_conn.delete_private_graph_endpoint.return_value = { + "status": "DELETING", + "vpcEndpointId": ENDPOINT_ID, + "vpcId": VPC_ID, + } + + operator = NeptuneDeletePrivateGraphEndpointOperator( + task_id="test_task", + graph_identifier=GRAPH_ID, + vpc_id=VPC_ID, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "vpc_id": VPC_ID, + "endpoint_id": ENDPOINT_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } + class TestNeptuneDeleteGraphOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -704,6 +803,37 @@ def test_delete_graph_other_client_error(self, mock_conn): with pytest.raises(NeptuneGraphDeletionFailedError): operator.execute(None) + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn): + mock_conn.delete_graph.return_value = { + "id": GRAPH_ID, + "name": GRAPH_NAME, + "status": "DELETING", + } + + operator = NeptuneDeleteGraphOperator( + task_id="test_task", + graph_id=GRAPH_ID, + skip_snapshot=True, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } + class TestNeptuneCreateGraphWithImportOperator: IMPORT_TASK_ID = "import-task-12345" @@ -971,6 +1101,69 @@ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn): assert exc_info.value.method_name == "defer_wait_for_task" assert exc_info.value.kwargs == {"import_task_id": self.IMPORT_TASK_ID} + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn): + mock_conn.create_graph_using_import_task.return_value = { + "graphId": GRAPH_ID, + "taskId": self.IMPORT_TASK_ID, + "status": "IMPORTING", + } + + operator = NeptuneCreateGraphWithImportOperator( + task_id="test_task", + graph_name=GRAPH_NAME, + vector_search_config={"dimension": 128}, + source=SOURCE_S3_URI, + role_arn=ROLE_ARN, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } + + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_defer_wait_for_task_trigger_receives_hook_configuration(self, mock_conn): + operator = NeptuneCreateGraphWithImportOperator( + task_id="test_task", + graph_name=GRAPH_NAME, + vector_search_config={"dimension": 128}, + source=SOURCE_S3_URI, + role_arn=ROLE_ARN, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.defer_wait_for_task( + import_task_id=self.IMPORT_TASK_ID, + context=None, + event={"status": "success"}, + ) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "import_task_id": self.IMPORT_TASK_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } + TASK_ID = "import-task-id-12345" @@ -1184,6 +1377,39 @@ def test_execute_complete_success(self): assert result == {"graph_id": GRAPH_ID, "import_task_id": TASK_ID} + @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneImportTaskLink.persist") + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_persist): + mock_conn.start_import_task.return_value = { + "taskId": TASK_ID, + "graphId": GRAPH_ID, + "status": "IMPORTING", + } + + operator = NeptuneStartImportTaskOperator( + task_id="test_task", + graph_identifier=GRAPH_ID, + role_arn=ROLE_ARN, + source=SOURCE_S3_URI, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "import_task_id": TASK_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } + class TestNeptuneCancelImportTaskOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -1275,3 +1501,33 @@ def test_execute_complete_success(self): result = operator.execute_complete(None, event) assert result == {"import_task_id": TASK_ID} + + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn): + mock_conn.cancel_import_task.return_value = { + "taskId": TASK_ID, + "graphId": GRAPH_ID, + "status": "CANCELLING", + } + + operator = NeptuneCancelImportTaskOperator( + task_id="test_task", + import_task_id=TASK_ID, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "task_identifier": TASK_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py index c6b114057eba9..d1da019247ebc 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py @@ -22,7 +22,7 @@ from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook from airflow.providers.amazon.aws.sensors.mwaa import MwaaDagRunSensor, MwaaTaskSensor -from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred from airflow.utils.state import DagRunState, TaskInstanceState SENSOR_DAG_RUN_KWARGS = { @@ -46,6 +46,10 @@ "max_retries": 100, } +REGION_NAME = "eu-west-2" +VERIFY = False +BOTOCORE_CONFIG = {"read_timeout": 42} + SENSOR_STATE_KWARGS = { "success_states": ["a", "b"], "failure_states": ["c", "d"], @@ -109,6 +113,32 @@ def test_execute_complete_success(self): success_event = {"status": "success", "dag_run_id": "test_run"} sensor.execute_complete({}, success_event) # should not raise + def test_deferred_trigger_receives_hook_configuration(self): + sensor = MwaaDagRunSensor( + **{**SENSOR_DAG_RUN_KWARGS, "deferrable": True}, + **SENSOR_STATE_KWARGS, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + sensor.execute({}) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": SENSOR_DAG_RUN_KWARGS["poke_interval"], + "waiter_max_attempts": SENSOR_DAG_RUN_KWARGS["max_retries"], + "aws_conn_id": "aws_default", + "external_env_name": SENSOR_DAG_RUN_KWARGS["external_env_name"], + "external_dag_id": SENSOR_DAG_RUN_KWARGS["external_dag_id"], + "external_dag_run_id": SENSOR_DAG_RUN_KWARGS["external_dag_run_id"], + "success_states": set(SENSOR_STATE_KWARGS["success_states"]), + "failure_states": set(SENSOR_STATE_KWARGS["failure_states"]), + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } + class TestMwaaTaskSuccessSensor: def test_init_success(self): @@ -159,3 +189,30 @@ def test_execute_complete_success(self): sensor = MwaaTaskSensor(**SENSOR_TASK_KWARGS, **SENSOR_STATE_KWARGS) success_event = {"status": "success", "task_id": "test_task"} sensor.execute_complete({}, success_event) # should not raise + + def test_deferred_trigger_receives_hook_configuration(self): + sensor = MwaaTaskSensor( + **SENSOR_TASK_KWARGS, + **SENSOR_STATE_KWARGS, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + sensor.execute({}) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": SENSOR_TASK_KWARGS["poke_interval"], + "waiter_max_attempts": SENSOR_TASK_KWARGS["max_retries"], + "aws_conn_id": "aws_default", + "external_env_name": SENSOR_TASK_KWARGS["external_env_name"], + "external_dag_id": SENSOR_TASK_KWARGS["external_dag_id"], + "external_dag_run_id": SENSOR_TASK_KWARGS["external_dag_run_id"], + "external_task_id": SENSOR_TASK_KWARGS["external_task_id"], + "success_states": set(SENSOR_STATE_KWARGS["success_states"]), + "failure_states": set(SENSOR_STATE_KWARGS["failure_states"]), + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py index 4b714d1e90dda..a8472776b4cbe 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py @@ -23,8 +23,12 @@ from airflow.providers.amazon.aws.hooks.ssm import SsmHook from airflow.providers.amazon.aws.sensors.ssm import SsmRunCommandCompletedSensor +from airflow.providers.common.compat.sdk import TaskDeferred COMMAND_ID = "123e4567-e89b-12d3-a456-426614174000" +REGION_NAME = "eu-west-2" +VERIFY = False +BOTOCORE_CONFIG = {"read_timeout": 42} @pytest.fixture @@ -144,3 +148,26 @@ def test_sensor_passes_parameter_to_trigger(self, mock_trigger_class): assert call_kwargs["command_id"] == COMMAND_ID assert call_kwargs["fail_on_nonzero_exit"] is False + + def test_deferred_trigger_receives_hook_configuration(self): + sensor = self.SENSOR( + **self.default_op_kwarg, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + sensor.execute({}) + + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": self.default_op_kwarg["poke_interval"], + "waiter_max_attempts": self.default_op_kwarg["max_retries"], + "aws_conn_id": "aws_default", + "command_id": COMMAND_ID, + "fail_on_nonzero_exit": True, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } diff --git a/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py new file mode 100644 index 0000000000000..8181274f04038 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py @@ -0,0 +1,286 @@ +# +# 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 ast +import importlib +import inspect +import pkgutil +from collections.abc import Iterator +from pathlib import Path + +import pytest + +import airflow.providers.amazon.aws as aws_module +import airflow.providers.amazon.aws.triggers as triggers_module +from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger + +AWS_ROOT = Path(inspect.getfile(aws_module)).parent +HOOK_CONFIGURATION = ("region_name", "verify", "botocore_config") + +# A deferrable task builds its hook twice: once in the worker, once in the triggerer. Unless the +# operator hands its hook configuration to the trigger, the triggerer silently falls back to boto3 +# defaults -- a different region, different SSL verification, different timeouts. +UNCONFIGURABLE_TRIGGERS = frozenset( + { + # Not an AwsBaseWaiterTrigger: its hook is addressed by execution name, and takes no + # connection parameters at all. + "SageMakerNotebookJobTrigger", + # A KubernetesPodTrigger; it reaches the pod through a kubeconfig, not a boto3 client. + "EksPodTrigger", + } +) + +# Sites whose trigger is built elsewhere and only referenced here, so the class cannot be read off +# the call. Kept explicit so that a new unreadable site fails the suite instead of being skipped. +UNREADABLE_DEFER_SITES = frozenset({("operators/eks.py", "trigger")}) + +# Services carved out as Contributors Workshop tasks, so their triggers are still unmigrated. Each +# entry is one self-contained contribution: widen the trigger's __init__, set aws_hook_class, pass +# the parameters at the call site, then delete the entry here. The test asserts an entry is still +# needed, so the allowlist cannot outlive the work it tracks. +PENDING_MIGRATION = frozenset( + { + ("sensors/batch.py", "BatchJobTrigger"), + ("sensors/opensearch_serverless.py", "OpenSearchServerlessCollectionActiveTrigger"), + } +) + + +def read_trigger_name(call: ast.Call) -> str | None: + """The trigger class a construction names, or ``None`` when the callee cannot be read.""" + if isinstance(call.func, ast.Name): + return call.func.id + if isinstance(call.func, ast.Attribute): + return call.func.attr + return None + + +def trigger_constructions(expr: ast.expr) -> list[tuple[ast.Call, str]] | None: + """ + Resolve a ``trigger=`` expression to the constructions it can evaluate to, each with its name. + + ``None`` means the expression cannot be read statically. Returning that rather than an empty + list is what keeps a site from disappearing: a bare reference, a subscript, or a conditional + with one unreadable branch all have to be acknowledged in ``UNREADABLE_DEFER_SITES`` instead of + quietly contributing nothing to the sweep. + """ + if isinstance(expr, ast.Call): + # A construction whose callee cannot be named is no more readable than a bare reference: + # the allowlists key on the class name, so an unnamed one could never match them. + name = read_trigger_name(expr) + return [(expr, name)] if name is not None else None + if isinstance(expr, ast.IfExp): + constructions: list[tuple[ast.Call, str]] = [] + for branch in (expr.body, expr.orelse): + resolved = trigger_constructions(branch) + if resolved is None: + return None + constructions.extend(resolved) + return constructions + return None + + +def walk_defer_sites() -> Iterator[tuple[Path, ast.expr]]: + """Yield the ``trigger=`` expression of every ``self.defer(...)`` in the provider.""" + # Every file, not just operators/ and sensors/: ``defer`` is a BaseOperator method, so a site + # can appear anywhere, and a directory filter would drop a nested subpackage without saying so. + for path in sorted(AWS_ROOT.rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text())): + if not isinstance(node, ast.Call): + continue + func = node.func + if not ( + isinstance(func, ast.Attribute) + and func.attr == "defer" + and isinstance(func.value, ast.Name) + and func.value.id == "self" + ): + continue + trigger = next((kw.value for kw in node.keywords if kw.arg == "trigger"), None) + if trigger is not None: + yield path, trigger + + +def find_defer_sites() -> list[tuple[str, int, str, list[str]]]: + """Collect every ``self.defer(trigger=SomeTrigger(...))`` in the provider.""" + sites: list[tuple[str, int, str, list[str]]] = [] + for path, trigger in walk_defer_sites(): + # The trigger may be built inline, or picked between in a conditional expression, so take + # every construction the expression can yield rather than assuming a single call. + for call, name in trigger_constructions(trigger) or (): + if name in UNCONFIGURABLE_TRIGGERS: + continue + passed = {kw.arg for kw in call.keywords if kw.arg} + sites.append( + ( + path.relative_to(AWS_ROOT).as_posix(), + call.lineno, + name, + [p for p in HOOK_CONFIGURATION if p not in passed], + ) + ) + return sites + + +def find_unreadable_defer_sites() -> set[tuple[str, str]]: + """Defer sites whose trigger expression cannot be resolved to the constructions it yields.""" + return { + (path.relative_to(AWS_ROOT).as_posix(), ast.unparse(trigger)) + for path, trigger in walk_defer_sites() + if trigger_constructions(trigger) is None + } + + +DEFER_SITES = find_defer_sites() + + +def test_defer_sites_are_discovered(): + assert DEFER_SITES, f"no self.defer(trigger=...) calls found under {AWS_ROOT}" + + +def test_no_defer_site_escapes_the_check(): + """A defer site whose trigger cannot be read statically must be acknowledged, not skipped.""" + assert find_unreadable_defer_sites() == UNREADABLE_DEFER_SITES + + +@pytest.mark.parametrize( + ("expression", "expected"), + [ + pytest.param("SomeTrigger(x=1)", 1, id="call"), + pytest.param("A() if flag else B()", 2, id="conditional-both-readable"), + pytest.param("trigger", None, id="bare-name"), + pytest.param("self._trigger", None, id="attribute"), + pytest.param("triggers[kind]", None, id="subscript"), + pytest.param("A() if flag else self._trigger", None, id="conditional-one-unreadable"), + pytest.param("TRIGGERS[kind](x=1)", None, id="unnameable-callee"), + pytest.param("module.SomeTrigger(x=1)", 1, id="module-qualified-callee"), + ], +) +def test_unreadable_trigger_expressions_resolve_to_none(expression, expected): + """Anything the sweep cannot resolve must report None so the site is forced onto the allowlist.""" + constructions = trigger_constructions(ast.parse(expression, mode="eval").body) + + assert (constructions if constructions is None else len(constructions)) == expected + + +@pytest.mark.parametrize( + ("source", "line", "trigger", "missing"), + DEFER_SITES, + ids=[f"{source}:{line}" for source, line, _, _ in DEFER_SITES], +) +def test_deferred_trigger_receives_hook_configuration(source, line, trigger, missing): + if (source, trigger) in PENDING_MIGRATION: + assert missing, ( + f"{source}:{line} now passes its hook configuration to {trigger}. " + f"Drop it from PENDING_MIGRATION so the site stays covered." + ) + pytest.skip(f"{source} is a Contributors Workshop task; see PENDING_MIGRATION") + + assert not missing, ( + f"{source}:{line} defers to {trigger} without passing {', '.join(missing)}. " + f"The triggerer builds its own hook, so anything not passed here is lost." + ) + + +def find_waiter_triggers() -> list[type[AwsBaseWaiterTrigger]]: + """Import every trigger module, then walk the subclass tree.""" + for module in pkgutil.iter_modules(triggers_module.__path__): + importlib.import_module(f"{triggers_module.__name__}.{module.name}") + + found: set[type[AwsBaseWaiterTrigger]] = set() + pending = [AwsBaseWaiterTrigger] + while pending: + for subclass in pending.pop().__subclasses__(): + if subclass not in found: + found.add(subclass) + pending.append(subclass) + return sorted(found, key=lambda cls: cls.__name__) + + +@pytest.mark.parametrize( + "trigger_class", + find_waiter_triggers(), + ids=lambda cls: cls.__name__, +) +def test_waiter_trigger_can_build_a_hook(trigger_class): + """The hook a trigger names must accept what ``_hook_parameters`` will pass it.""" + if trigger_class.hook is not AwsBaseWaiterTrigger.hook: + pytest.skip(f"{trigger_class.__name__} builds its hook by hand") + inspect.signature(trigger_class.aws_hook_class).bind_partial( + aws_conn_id=None, region_name=None, verify=None, config=None + ) + + +# A trigger may build a second hook by hand for a side channel -- streaming CloudWatch logs, most +# often -- alongside the one ``aws_hook_class`` gives it. That hook talks to AWS too, so it needs +# the same configuration; a trigger whose job client verifies TLS while its log client does not is +# the same bug in miniature. +HAND_BUILT_HOOK_EXCEPTIONS = frozenset( + { + # Addressed by execution name; takes no connection parameters at all. + ("sagemaker_unified_studio.py", "SageMakerNotebookHook"), + # EksPodOperator is a KubernetesPodOperator: it carries no verify or botocore_config to pass. + ("eks.py", "EksHook"), + # Contributors Workshop task; see PENDING_MIGRATION. + ("opensearch_serverless.py", "OpenSearchServerlessHook"), + } +) + + +def find_hand_built_hooks() -> list[tuple[str, int, str, list[str]]]: + """Collect every hook constructed directly inside a trigger module.""" + sites = [] + for path in sorted((AWS_ROOT / "triggers").rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text())): + if not isinstance(node, ast.Call): + continue + # Match a module-qualified ``module.SomeHook(...)`` as well as a bare name. + hook = node.func.attr if isinstance(node.func, ast.Attribute) else getattr(node.func, "id", "") + if not hook.endswith("Hook"): + continue + passed = {keyword.arg for keyword in node.keywords if keyword.arg} + # AwsGenericHook names the botocore config ``config``. + if "config" in passed: + passed.add("botocore_config") + missing = [name for name in HOOK_CONFIGURATION if name not in passed] + sites.append((path.name, node.lineno, hook, missing)) + return sites + + +HAND_BUILT_HOOKS = find_hand_built_hooks() + + +@pytest.mark.parametrize( + ("source", "line", "hook", "missing"), + HAND_BUILT_HOOKS, + ids=[f"{source}:{line}" for source, line, _, _ in HAND_BUILT_HOOKS], +) +def test_hand_built_trigger_hook_receives_configuration(source, line, hook, missing): + """A hook a trigger builds itself must carry the same configuration as its main hook.""" + if (source, hook) in HAND_BUILT_HOOK_EXCEPTIONS: + assert missing, ( + f"{source}:{line} now configures {hook}. " + f"Drop it from HAND_BUILT_HOOK_EXCEPTIONS so the site stays covered." + ) + pytest.skip(f"{source} builds {hook} with nothing to configure") + + assert not missing, ( + f"{source}:{line} builds {hook} without {', '.join(missing)}. " + f"It reaches AWS with boto3 defaults while the trigger's own hook does not." + ) diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py b/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py index 0c2ea69149a06..60454ee45215b 100644 --- a/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py +++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py @@ -23,10 +23,10 @@ import pytest from airflow.exceptions import AirflowException +from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook from airflow.triggers.base import TriggerEvent @@ -156,3 +156,37 @@ def test_event_from_exception(self): "message": "AWS Glue job failed.", "hello": "world", } + + +class TestAwsBaseWaiterTriggerSubclassValidation: + """``hook()`` runs in the triggerer, so a subclass that cannot build one must fail on import.""" + + def test_subclass_with_neither_hook_class_nor_hook_is_rejected(self): + with pytest.raises(AttributeError, match="aws_hook_class' should be set"): + + class MissingBoth(AwsBaseWaiterTrigger): + pass + + def test_subclass_whose_hook_class_is_not_a_hook_is_rejected(self): + with pytest.raises(AttributeError, match="not a subclass of AwsGenericHook"): + + class NotAHook(AwsBaseWaiterTrigger): + aws_hook_class = str + + def test_subclass_declaring_a_hook_class_is_accepted(self): + class Declared(AwsBaseWaiterTrigger): + aws_hook_class = AwsGenericHook + + assert Declared.aws_hook_class is AwsGenericHook + + def test_subclass_inheriting_a_hook_override_is_accepted(self): + """An override reached through an intermediate base still counts.""" + + class Intermediate(AwsBaseWaiterTrigger): + def hook(self) -> AwsGenericHook: + return AsyncMock() + + class Leaf(Intermediate): + pass + + assert Leaf.hook is Intermediate.hook