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 71fa6bf1dd110..59a49707d538b 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..99196b01c551d 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", 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 aab675201497c..5faea9d8885c6 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", @@ -1683,6 +1710,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, ), @@ -1706,6 +1736,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, ), @@ -1736,6 +1769,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, ), @@ -1829,6 +1865,9 @@ def execute(self, context: Context) -> 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 c488965ddde01..c1d688f4e40b9 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py @@ -364,6 +364,8 @@ def execute(self, context: Context) -> str | None: run_id=job_run_id, verbose=self.verbose, 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, @@ -754,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", ) @@ -902,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 5d00918ec46fe..025440187d18a 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 f9a632893882e..b6fe56565a782 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/emr.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/emr.py @@ -326,6 +326,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, ) @@ -334,6 +337,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", @@ -523,6 +529,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", @@ -653,6 +662,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 c5c710e9e8fb3..d90090a627844 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/glue.py @@ -102,6 +102,8 @@ def execute(self, context: Context) -> Any: run_id=self.run_id, verbose=self.verbose, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, region_name=self.region_name, @@ -210,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", ) @@ -330,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 999b5e5bfd128..49e105cfcef83 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py @@ -17,7 +17,6 @@ from __future__ import annotations -from abc import abstractmethod from collections.abc import AsyncIterator from typing import TYPE_CHECKING, Any @@ -34,7 +33,11 @@ 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 +70,9 @@ 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__( self, *, @@ -137,9 +143,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) async def run(self) -> AsyncIterator[TriggerEvent]: hook = self.hook() 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 abf0bb4c9db82..d7a95ddf85d19 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..1952bcb987fd5 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,14 @@ 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 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 +397,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 +412,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 +426,14 @@ 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 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 +442,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 +457,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 +476,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 +492,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 +511,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 +530,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 +546,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 +561,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 48ddb4b0c3197..3d12ed074b324 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 EmrServerlessStartJobTrigger(AwsBaseWaiterTrigger): """ @@ -381,8 +465,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, @@ -391,6 +482,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={ @@ -408,14 +502,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 @@ -557,14 +651,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}, @@ -578,11 +682,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): """ @@ -592,14 +696,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__( @@ -614,11 +728,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..097ecd3ab968c 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(): @@ -322,14 +313,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 +344,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 +358,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 +389,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_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..d887c51ec93b2 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py @@ -0,0 +1,193 @@ +# +# 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 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( + { + ("operators/sagemaker_unified_studio_notebook.py", "SageMakerUnifiedStudioNotebookTrigger"), + ("sensors/batch.py", "BatchJobTrigger"), + ("sensors/opensearch_serverless.py", "OpenSearchServerlessCollectionActiveTrigger"), + } +) + + +def trigger_constructions(expr: ast.expr) -> list[ast.Call]: + """Resolve a ``trigger=`` expression to the constructions it can evaluate to.""" + if isinstance(expr, ast.Call): + return [expr] + if isinstance(expr, ast.IfExp): + return trigger_constructions(expr.body) + trigger_constructions(expr.orelse) + return [] + + +def find_defer_sites() -> list[tuple[str, int, str, list[str]]]: + """Collect every ``self.defer(trigger=SomeTrigger(...))`` in the provider.""" + sites = [] + for path in sorted(AWS_ROOT.rglob("*.py")): + if path.parent.name not in ("operators", "sensors"): + continue + 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 None: + continue + # 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 in trigger_constructions(trigger): + name = ( + call.func.attr if isinstance(call.func, ast.Attribute) else getattr(call.func, "id", "") + ) + if name in UNCONFIGURABLE_TRIGGERS: + continue + passed = {kw.arg for kw in call.keywords if kw.arg} + sites.append( + ( + str(path.relative_to(AWS_ROOT)), + 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 is a bare reference, so its class cannot be read statically.""" + unreadable = set() + for path in sorted(AWS_ROOT.rglob("*.py")): + if path.parent.name not in ("operators", "sensors"): + continue + 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 isinstance(trigger, ast.Name): + unreadable.add((str(path.relative_to(AWS_ROOT)), trigger.id)) + return unreadable + + +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( + ("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): + """Every waiter trigger must declare ``aws_hook_class`` or provide its own ``hook()``.""" + assert hasattr(trigger_class, "aws_hook_class") or "hook" in vars(trigger_class), ( + f"{trigger_class.__name__} sets neither aws_hook_class nor hook(); " + f"building its hook would fail at runtime." + )