From 15002fe7b86c7f29df0749a5fa02447e88c35e50 Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Wed, 26 Aug 2026 14:52:31 +0530 Subject: [PATCH 1/3] Use the operator's AWS settings for deferred Neptune, MWAA, SSM tasks An AwsBaseOperator/AwsBaseSensor subclass resolves region_name, verify and botocore_config in __init__, but did not hand them to the trigger it defers to. The trigger builds its own hook, so the deferred half of the task reached AWS with the default region, SSL verification silently re-enabled, and any custom botocore timeouts or retries discarded. The triggers already accept all three, so only the call sites were missing. --- .../providers/amazon/aws/operators/mwaa.py | 3 + .../amazon/aws/operators/neptune_analytics.py | 24 ++ .../providers/amazon/aws/sensors/mwaa.py | 6 + .../providers/amazon/aws/sensors/ssm.py | 3 + .../unit/amazon/aws/operators/test_mwaa.py | 21 ++ .../aws/operators/test_neptune_analytics.py | 205 ++++++++++++++++++ .../unit/amazon/aws/sensors/test_mwaa.py | 38 +++- .../tests/unit/amazon/aws/sensors/test_ssm.py | 20 ++ 8 files changed, 319 insertions(+), 1 deletion(-) 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..239906e5e9358 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 @@ -173,6 +173,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, @@ -317,6 +320,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, @@ -420,6 +426,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, @@ -517,6 +526,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, @@ -729,6 +741,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 +788,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}, @@ -914,6 +932,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", ) @@ -1002,6 +1023,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..e2232f77384fd 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,20 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.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..f5aea2c2fc95a 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,29 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneCreatePrivateGraphEndpointOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -417,6 +443,33 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneDeletePrivateGraphEndpointOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -548,6 +601,31 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneDeleteGraphOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -704,6 +782,31 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneCreateGraphWithImportOperator: IMPORT_TASK_ID = "import-task-12345" @@ -971,6 +1074,57 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + TASK_ID = "import-task-id-12345" @@ -1184,6 +1338,33 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneCancelImportTaskOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -1275,3 +1456,27 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.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..000cc710a887f 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,22 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestMwaaTaskSuccessSensor: def test_init_success(self): @@ -159,3 +179,19 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.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..9b4bdf80072c1 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,19 @@ 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.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG From ef8ed9ead67bd2525bf86dc20c5a8de638da433b Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Thu, 27 Aug 2026 14:24:14 +0530 Subject: [PATCH 2/3] Document the verify parameter for Neptune Analytics operators Every Neptune Analytics operator accepts `verify` through the shared AWS base class, but none of the seven class docstrings mentioned it, so the rendered provider docs gave users no way to discover it. Two of those docstrings even carried a stray blank line where the entry belonged. The deferral tests now compare the trigger's serialized payload rather than its attributes. Serialization is what actually crosses into the triggerer process, and it passes values through `prune_dict`, so an attribute-level assertion can pass while the setting is silently dropped on the way there. This matches the assertion style already used for the Neptune cluster operators. --- .../amazon/aws/operators/neptune_analytics.py | 16 ++- .../unit/amazon/aws/operators/test_mwaa.py | 16 ++- .../aws/operators/test_neptune_analytics.py | 99 ++++++++++++++----- .../unit/amazon/aws/sensors/test_mwaa.py | 33 +++++-- .../tests/unit/amazon/aws/sensors/test_ssm.py | 13 ++- 5 files changed, 139 insertions(+), 38 deletions(-) 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 239906e5e9358..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 @@ -233,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 @@ -379,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 @@ -482,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 @@ -594,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 @@ -840,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 @@ -983,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 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 e2232f77384fd..4e91ab1e1b472 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py @@ -133,6 +133,16 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_hook): with pytest.raises(TaskDeferred) as exc_info: op.execute({}) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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 f5aea2c2fc95a..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 @@ -241,9 +241,15 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_pers with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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: @@ -466,9 +472,16 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_pers with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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: @@ -622,9 +635,17 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn): with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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: @@ -803,9 +824,15 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn): with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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: @@ -1097,9 +1124,15 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn): with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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): @@ -1121,9 +1154,15 @@ def test_defer_wait_for_task_trigger_receives_hook_configuration(self, mock_conn event={"status": "success"}, ) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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" @@ -1361,9 +1400,15 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_pers with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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: @@ -1477,6 +1522,12 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn): with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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 000cc710a887f..d1da019247ebc 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py @@ -125,9 +125,19 @@ def test_deferred_trigger_receives_hook_configuration(self): with pytest.raises(TaskDeferred) as exc_info: sensor.execute({}) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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: @@ -192,6 +202,17 @@ def test_deferred_trigger_receives_hook_configuration(self): with pytest.raises(TaskDeferred) as exc_info: sensor.execute({}) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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 9b4bdf80072c1..a8472776b4cbe 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py @@ -161,6 +161,13 @@ def test_deferred_trigger_receives_hook_configuration(self): with pytest.raises(TaskDeferred) as exc_info: sensor.execute({}) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + 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, + } From d7f8a6ec53309f3b916e8e57b1bcfa8352c45715 Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Thu, 27 Aug 2026 23:47:34 +0530 Subject: [PATCH 3/3] Build deferred AWS hooks from the operator's own settings An AWS operator always carries region_name, verify and botocore_config, but on deferral the trigger builds its own hook. Most triggers accepted none of those parameters and constructed the hook from aws_conn_id alone, so the triggerer silently fell back to boto3 defaults: a different region, default SSL verification, and none of the configured timeouts or retry policy. The task changed behaviour purely by virtue of deferring, and did so without any error. Fixing this service by service would have meant editing every trigger signature as well as every call site, so the hook is now built in one place from the parameters the base trigger already serializes. Subclasses name the hook they need instead of constructing it, which is the same arrangement the operators use. The accompanying invariant test walks every defer site in the provider and fails if one does not hand its hook configuration to the trigger, so an operator added later cannot reintroduce the gap unnoticed. Three services are deliberately left for the Contributors Workshop and are named in the test's allowlist rather than skipped silently. --- .../providers/amazon/aws/operators/bedrock.py | 21 ++ .../amazon/aws/operators/comprehend.py | 6 + .../providers/amazon/aws/operators/dms.py | 21 ++ .../providers/amazon/aws/operators/ecs.py | 4 + .../providers/amazon/aws/operators/eks.py | 18 ++ .../providers/amazon/aws/operators/emr.py | 39 ++++ .../providers/amazon/aws/operators/glue.py | 8 + .../providers/amazon/aws/operators/rds.py | 8 + .../providers/amazon/aws/sensors/bedrock.py | 15 ++ .../amazon/aws/sensors/comprehend.py | 6 + .../providers/amazon/aws/sensors/emr.py | 12 ++ .../providers/amazon/aws/sensors/glue.py | 8 + .../providers/amazon/aws/triggers/base.py | 29 ++- .../providers/amazon/aws/triggers/bedrock.py | 139 ++++++++++--- .../amazon/aws/triggers/comprehend.py | 37 +++- .../providers/amazon/aws/triggers/dms.py | 118 ++++++----- .../providers/amazon/aws/triggers/ecs.py | 31 ++- .../providers/amazon/aws/triggers/eks.py | 82 ++++++-- .../providers/amazon/aws/triggers/emr.py | 180 +++++++++++++--- .../providers/amazon/aws/triggers/glue.py | 47 +++-- .../providers/amazon/aws/triggers/rds.py | 45 ++-- .../aws/test_deferred_hook_configuration.py | 193 ++++++++++++++++++ 22 files changed, 884 insertions(+), 183 deletions(-) create mode 100644 providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py 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/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/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/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." + )