Use the operator's AWS settings for deferred Neptune, MWAA, SSM tasks

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",
             )

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",
             )

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",
             )

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",

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,
+        }