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