From 0c3c18b9d3057331053c76bf86260868df438656 Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Wed, 26 Aug 2026 14:52:31 +0530 Subject: [PATCH 1/9] 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 3c43a3eb389709f482ef463fd00b96475f560964 Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Thu, 27 Aug 2026 14:24:14 +0530 Subject: [PATCH 2/9] 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 414e5d896163a68eafe9ffccb10dbab16eda31fb Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Thu, 27 Aug 2026 23:47:34 +0530 Subject: [PATCH 3/9] 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 | 9 +- .../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 | 84 ++++-- .../providers/amazon/aws/triggers/emr.py | 180 ++++++++++--- .../providers/amazon/aws/triggers/glue.py | 52 ++-- .../providers/amazon/aws/triggers/rds.py | 45 +++- .../unit/amazon/aws/operators/test_ecs.py | 37 ++- .../aws/test_deferred_hook_configuration.py | 251 ++++++++++++++++++ 23 files changed, 988 insertions(+), 186 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..be6c171a56a3a 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py @@ -141,6 +141,8 @@ def execute(self, context: Context): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, ), method_name="_complete_exec_with_cluster_desc", @@ -218,6 +220,8 @@ def execute(self, context: Context): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, ), method_name="_complete_exec_with_cluster_desc", @@ -622,7 +626,10 @@ def execute_complete(self, context: Context, event: dict[str, Any] | None = None if self._aws_logs_enabled(): # same behavior as non-deferrable mode, return last line of logs of the task. logs_client = AwsLogsHook( - aws_conn_id=self.aws_conn_id, region_name=self.resolve_awslogs_region() + aws_conn_id=self.aws_conn_id, + region_name=self.resolve_awslogs_region(), + verify=self.verify, + config=self.botocore_config, ).conn one_log = logs_client.get_log_events( logGroupName=self.awslogs_group, diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py index c2efc42f8a57b..e4fa2efde429d 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py @@ -359,6 +359,8 @@ def execute(self, context: Context): trigger=EksCreateClusterTrigger( cluster_name=self.cluster_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -411,6 +413,8 @@ def deferrable_create_cluster_next(self, context: Context, event: dict[str, Any] waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, force_delete_compute=False, ), @@ -445,6 +449,8 @@ def deferrable_create_cluster_next(self, context: Context, event: dict[str, Any] waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, ), method_name="execute_complete", @@ -456,6 +462,8 @@ def deferrable_create_cluster_next(self, context: Context, event: dict[str, Any] nodegroup_name=self.nodegroup_name, cluster_name=self.cluster_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -600,6 +608,8 @@ def execute(self, context: Context): cluster_name=self.cluster_name, nodegroup_name=self.nodegroup_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -714,6 +724,8 @@ def execute(self, context: Context): cluster_name=self.cluster_name, fargate_profile_name=self.fargate_profile_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, region_name=self.region_name, @@ -805,6 +817,8 @@ def execute(self, context: Context): waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, force_delete_compute=self.force_delete_compute, ), @@ -946,6 +960,8 @@ def execute(self, context: Context): cluster_name=self.cluster_name, nodegroup_name=self.nodegroup_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, region_name=self.region_name, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -1038,6 +1054,8 @@ def execute(self, context: Context): cluster_name=self.cluster_name, fargate_profile_name=self.fargate_profile_name, aws_conn_id=self.aws_conn_id, + verify=self.verify, + botocore_config=self.botocore_config, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, region_name=self.region_name, diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py index 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 cfb8575a752d0..136165752acfd 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 b952e27fbbaf2..840429388dbab 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) def _event_from_exception(self, error: AirflowException) -> TriggerEvent: return TriggerEvent( diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py index 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..1d142d9dbb645 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py @@ -29,12 +29,11 @@ from airflow.providers.cncf.kubernetes.triggers.pod import KubernetesPodTrigger from airflow.providers.common.compat.sdk import AirflowException from airflow.triggers.base import TriggerEvent +from airflow.utils.helpers import prune_dict if TYPE_CHECKING: from pendulum import DateTime - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - class EksCreateClusterTrigger(AwsBaseWaiterTrigger): """ @@ -48,8 +47,14 @@ class EksCreateClusterTrigger(AwsBaseWaiterTrigger): :param aws_conn_id: The Airflow connection used for AWS credentials. :param region_name: Which AWS region the connection should use. If this is None or empty then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -57,6 +62,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"cluster_name": cluster_name, "region_name": region_name}, @@ -70,11 +77,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - async def run(self): async with await self.hook().get_async_conn() as client: waiter = client.get_waiter(self.waiter_name) @@ -237,8 +243,14 @@ class EksDeleteClusterTrigger(AwsBaseWaiterTrigger): If this is None or empty then the default boto3 behaviour is used. :param force_delete_compute: If True, any nodegroups or fargate profiles associated with the cluster will be deleted before the cluster is deleted. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = EksHook + def __init__( self, cluster_name, @@ -247,12 +259,16 @@ def __init__( aws_conn_id: str | None, region_name: str | None, force_delete_compute: bool, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): self.cluster_name = cluster_name self.waiter_delay = waiter_delay self.waiter_max_attempts = waiter_max_attempts self.aws_conn_id = aws_conn_id self.region_name = region_name + self.verify = verify + self.botocore_config = botocore_config self.force_delete_compute = force_delete_compute def serialize(self) -> tuple[str, dict[str, Any]]: @@ -265,12 +281,10 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "aws_conn_id": self.aws_conn_id, "region_name": self.region_name, "force_delete_compute": self.force_delete_compute, + **prune_dict({"verify": self.verify, "botocore_config": self.botocore_config}), }, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - async def run(self): async with await self.hook().get_async_conn() as client: waiter = client.get_waiter("cluster_deleted") @@ -367,8 +381,15 @@ class EksCreateFargateProfileTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: Which AWS region the connection should use. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -377,6 +398,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"cluster_name": cluster_name, "fargate_profile_name": fargate_profile_name}, @@ -390,11 +413,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class EksDeleteFargateProfileTrigger(AwsBaseWaiterTrigger): """ @@ -405,8 +427,15 @@ class EksDeleteFargateProfileTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: Which AWS region the connection should use. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -415,6 +444,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"cluster_name": cluster_name, "fargate_profile_name": fargate_profile_name}, @@ -428,11 +459,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger): """ @@ -448,8 +478,14 @@ class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger): :param aws_conn_id: The Airflow connection used for AWS credentials. :param region_name: Which AWS region the connection should use. (templated) If this is None or empty then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -458,6 +494,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={ @@ -475,11 +513,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class EksDeleteNodegroupTrigger(AwsBaseWaiterTrigger): """ @@ -495,8 +532,14 @@ class EksDeleteNodegroupTrigger(AwsBaseWaiterTrigger): :param aws_conn_id: The Airflow connection used for AWS credentials. :param region_name: Which AWS region the connection should use. (templated) If this is None or empty then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = EksHook + def __init__( self, cluster_name: str, @@ -505,6 +548,8 @@ def __init__( waiter_max_attempts: int, aws_conn_id: str | None, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"cluster_name": cluster_name, "nodegroup_name": nodegroup_name}, @@ -518,7 +563,6 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - - def hook(self) -> AwsGenericHook: - return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py index 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..42c0e501decd8 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py @@ -20,13 +20,10 @@ import asyncio from collections.abc import AsyncIterator from functools import cached_property -from typing import TYPE_CHECKING, Any +from typing import Any from botocore.exceptions import ClientError -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - from airflow.providers.amazon.aws.hooks.glue import ( GlueDataQualityHook, GlueJobHook, @@ -55,6 +52,8 @@ class GlueJobCompleteTrigger(AwsBaseWaiterTrigger): :param botocore_config: Configuration dictionary (key-values) for botocore client. """ + aws_hook_class = GlueJobHook + def __init__( self, job_name: str, @@ -87,14 +86,6 @@ def __init__( self.run_id = run_id self.verbose = verbose - def hook(self) -> AwsGenericHook: - return GlueJobHook( - aws_conn_id=self.aws_conn_id, - region_name=self.region_name, - verify=self.verify, - config=self.botocore_config, - ) - async def run(self) -> AsyncIterator[TriggerEvent]: if not self.verbose: async for event in super().run(): @@ -105,7 +96,10 @@ async def run(self) -> AsyncIterator[TriggerEvent]: async with ( await hook.get_async_conn() as glue_client, await AwsLogsHook( - aws_conn_id=self.aws_conn_id, region_name=self.region_name + aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + config=self.botocore_config, ).get_async_conn() as logs_client, ): # Get log group names from job run metadata @@ -322,14 +316,24 @@ class GlueDataQualityRuleSetEvaluationRunCompleteTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60) :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75) :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = GlueDataQualityHook + def __init__( self, evaluation_run_id: str, waiter_delay: int = 60, waiter_max_attempts: int = 75, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"evaluation_run_id": evaluation_run_id}, @@ -343,11 +347,11 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return GlueDataQualityHook(aws_conn_id=self.aws_conn_id) - class GlueDataQualityRuleRecommendationRunCompleteTrigger(AwsBaseWaiterTrigger): """ @@ -357,14 +361,24 @@ class GlueDataQualityRuleRecommendationRunCompleteTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60) :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75) :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: The AWS region where the resources to watch are. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = GlueDataQualityHook + def __init__( self, recommendation_run_id: str, waiter_delay: int = 60, waiter_max_attempts: int = 75, aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ): super().__init__( serialized_fields={"recommendation_run_id": recommendation_run_id}, @@ -378,7 +392,7 @@ def __init__( waiter_delay=waiter_delay, waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - - def hook(self) -> AwsGenericHook: - return GlueDataQualityHook(aws_conn_id=self.aws_conn_id) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/rds.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/rds.py index dd60037c69787..0e1ed8439bd6c 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/rds.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/rds.py @@ -16,16 +16,12 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any from airflow.providers.amazon.aws.hooks.rds import RdsHook from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger from airflow.providers.amazon.aws.utils.rds import RdsDbType -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - - _waiter_arg = { RdsDbType.INSTANCE.value: "DBInstanceIdentifier", RdsDbType.CLUSTER.value: "DBClusterIdentifier", @@ -47,8 +43,14 @@ class RdsDbAvailableTrigger(AwsBaseWaiterTrigger): :param region_name: AWS region where the DB is located, if different from the default one. :param response: The response from the RdsHook, to be passed back to the operator. :param db_type: The type of DB: instance or cluster. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = RdsHook + def __init__( self, db_identifier: str, @@ -58,6 +60,8 @@ def __init__( response: dict[str, Any], db_type: RdsDbType | str, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: # allow passing enums for users, # but we can only rely on strings because (de-)serialization doesn't support enums @@ -83,11 +87,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class RdsDbDeletedTrigger(AwsBaseWaiterTrigger): """ @@ -100,8 +103,14 @@ class RdsDbDeletedTrigger(AwsBaseWaiterTrigger): :param region_name: AWS region where the DB is located, if different from the default one. :param response: The response from the RdsHook, to be passed back to the operator. :param db_type: The type of DB: instance or cluster. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = RdsHook + def __init__( self, db_identifier: str, @@ -111,6 +120,8 @@ def __init__( response: dict[str, Any], db_type: RdsDbType | str, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: # allow passing enums for users, # but we can only rely on strings because (de-)serialization doesn't support enums @@ -136,11 +147,10 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - def hook(self) -> AwsGenericHook: - return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) - class RdsDbStoppedTrigger(AwsBaseWaiterTrigger): """ @@ -153,8 +163,14 @@ class RdsDbStoppedTrigger(AwsBaseWaiterTrigger): :param region_name: AWS region where the DB is located, if different from the default one. :param response: The response from the RdsHook, to be passed back to the operator. :param db_type: The type of DB: instance or cluster. + :param verify: Whether or not to verify SSL certificates. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """ + aws_hook_class = RdsHook + def __init__( self, db_identifier: str, @@ -164,6 +180,8 @@ def __init__( response: dict[str, Any], db_type: RdsDbType | str, region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, ) -> None: # allow passing enums for users, # but we can only rely on strings because (de-)serialization doesn't support enums @@ -189,7 +207,6 @@ def __init__( waiter_max_attempts=waiter_max_attempts, aws_conn_id=aws_conn_id, region_name=region_name, + verify=verify, + botocore_config=botocore_config, ) - - def hook(self) -> AwsGenericHook: - return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py b/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py index 8a28069fd255e..137ab9c64c1e8 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py @@ -901,7 +901,42 @@ def test_execute_complete_uses_awslogs_region(self, check_mock, logs_hook_mock): assert result == "Log output" check_mock.assert_called_once_with() - logs_hook_mock.assert_called_once_with(aws_conn_id=self.ecs.aws_conn_id, region_name="logs-region") + logs_hook_mock.assert_called_once_with( + aws_conn_id=self.ecs.aws_conn_id, + region_name="logs-region", + verify=self.ecs.verify, + config=self.ecs.botocore_config, + ) + + @mock.patch("airflow.providers.amazon.aws.operators.ecs.AwsLogsHook") + @mock.patch.object(EcsRunTaskOperator, "_check_success_task") + def test_execute_complete_log_hook_uses_operator_aws_configuration(self, check_mock, logs_hook_mock): + botocore_config = {"read_timeout": 10} + self.set_up_operator( + awslogs_group="awslogs-group", + awslogs_region="logs-region", + awslogs_stream_prefix="prefix", + region_name="task-region", + verify="/path/to/ca-bundle.pem", + botocore_config=botocore_config, + ) + logs_hook_mock.return_value.conn.get_log_events.return_value = {"events": [{"message": "Log output"}]} + + self.ecs.execute_complete( + {}, + { + "status": "success", + "task_arn": f"arn:aws:ecs:us-east-1:012345678910:task/{TASK_ID}", + "cluster": "test_cluster", + }, + ) + + logs_hook_mock.assert_called_once_with( + aws_conn_id=self.ecs.aws_conn_id, + region_name="logs-region", + verify="/path/to/ca-bundle.pem", + config=botocore_config, + ) @mock.patch.object(EcsBaseOperator, "client") @mock.patch("airflow.providers.amazon.aws.utils.task_log_fetcher.AwsTaskLogFetcher") diff --git a/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py new file mode 100644 index 0000000000000..197cdffdcd373 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py @@ -0,0 +1,251 @@ +# +# 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( + { + ("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." + ) + + +# A trigger may build a second hook by hand for a side channel -- streaming CloudWatch logs, most +# often -- alongside the one ``aws_hook_class`` gives it. That hook talks to AWS too, so it needs +# the same configuration; a trigger whose job client verifies TLS while its log client does not is +# the same bug in miniature. +HAND_BUILT_HOOK_EXCEPTIONS = frozenset( + { + # Addressed by execution name; takes no connection parameters at all. + ("sagemaker_unified_studio.py", "SageMakerNotebookHook"), + # EksPodOperator is a KubernetesPodOperator: it carries no verify or botocore_config to pass. + ("eks.py", "EksHook"), + # Contributors Workshop task; see PENDING_MIGRATION. + ("opensearch_serverless.py", "OpenSearchServerlessHook"), + } +) + + +def find_hand_built_hooks() -> list[tuple[str, int, str, list[str]]]: + """Collect every hook constructed directly inside a trigger module.""" + sites = [] + for path in sorted((AWS_ROOT / "triggers").rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text())): + if not ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id.endswith("Hook") + ): + continue + passed = {keyword.arg for keyword in node.keywords if keyword.arg} + # AwsGenericHook names the botocore config ``config``. + if "config" in passed: + passed.add("botocore_config") + missing = [name for name in HOOK_CONFIGURATION if name not in passed] + sites.append((path.name, node.lineno, node.func.id, missing)) + return sites + + +HAND_BUILT_HOOKS = find_hand_built_hooks() + + +@pytest.mark.parametrize( + ("source", "line", "hook", "missing"), + HAND_BUILT_HOOKS, + ids=[f"{source}:{line}" for source, line, _, _ in HAND_BUILT_HOOKS], +) +def test_hand_built_trigger_hook_receives_configuration(source, line, hook, missing): + """A hook a trigger builds itself must carry the same configuration as its main hook.""" + if (source, hook) in HAND_BUILT_HOOK_EXCEPTIONS: + assert missing, ( + f"{source}:{line} now configures {hook}. " + f"Drop it from HAND_BUILT_HOOK_EXCEPTIONS so the site stays covered." + ) + pytest.skip(f"{source} builds {hook} with nothing to configure") + + assert not missing, ( + f"{source}:{line} builds {hook} without {', '.join(missing)}. " + f"It reaches AWS with boto3 defaults while the trigger's own hook does not." + ) From 91e1923d54f89bcd65ec5f20c962ea21ffb6dc12 Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Thu, 10 Sep 2026 08:21:14 +0530 Subject: [PATCH 4/9] Drop the duplicated Glue trigger arguments after merging main Merging main brought in #72557, which passes verify and botocore_config to GlueJobCompleteTrigger at the same two call sites this branch already widened. Git combined both insertions without reporting a conflict, leaving each call with the arguments repeated, which Python rejects at import time. --- .../amazon/src/airflow/providers/amazon/aws/operators/glue.py | 2 -- .../amazon/src/airflow/providers/amazon/aws/sensors/glue.py | 2 -- 2 files changed, 4 deletions(-) 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 b3583704de389..ca2ff750d300e 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py @@ -364,8 +364,6 @@ 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, 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 226f915436c2e..58190131c190f 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/glue.py @@ -102,8 +102,6 @@ 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, From 6eed096804fe7a29ddf3840027cf8281dc12596f Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Thu, 10 Sep 2026 08:51:51 +0530 Subject: [PATCH 5/9] Fail fast when an AWS trigger cannot build its hook hook() is only reached from run(), which executes in the triggerer, so a subclass declaring neither aws_hook_class nor its own hook() would defer successfully and fail later, out of sight of the task that deferred. The operator side already gets this guarantee from validate_attributes. Checked on class creation rather than in __init__ because subclasses such as EksDeleteClusterTrigger never call super().__init__(). Also records the behaviour change in the provider changelog, since the triggerer now uses the operator's region, SSL verification and botocore configuration rather than falling back to boto3 defaults. --- providers/amazon/docs/changelog.rst | 9 +++++ .../providers/amazon/aws/triggers/base.py | 31 +++++++++++++--- .../aws/test_deferred_hook_configuration.py | 18 +++++----- .../unit/amazon/aws/triggers/test_base.py | 36 ++++++++++++++++++- 4 files changed, 80 insertions(+), 14 deletions(-) diff --git a/providers/amazon/docs/changelog.rst b/providers/amazon/docs/changelog.rst index ca7ffc5acaa88..47e08b2443f27 100644 --- a/providers/amazon/docs/changelog.rst +++ b/providers/amazon/docs/changelog.rst @@ -26,6 +26,15 @@ Changelog --------- +.. warning:: + Deferrable AWS operators and sensors now hand ``region_name``, ``verify`` and ``botocore_config`` + to the trigger they defer to, so the triggerer builds its hook from the operator's settings + instead of falling back to boto3 defaults. Deployments where the triggerer happened to work + *because* of those defaults will see it change: it now uses the operator's region rather than the + triggerer host's ``AWS_DEFAULT_REGION``, and it applies the operator's SSL verification and + botocore configuration, which previously never reached it. Set these explicitly on the operator + if the deferred half needs to differ from the synchronous half. + .. warning:: The default waiter timeout of ``ComprehendCreateDocumentClassifierOperator`` was raised from 20 minutes (``waiter_max_attempts=20``) to 60 minutes (``waiter_max_attempts=60``), because 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 840429388dbab..97f5b2c375014 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py @@ -18,16 +18,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.utils.waiter_with_logging import async_wait from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.utils.helpers import prune_dict -if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook - class AwsBaseWaiterTrigger(BaseTrigger): """ @@ -73,6 +71,31 @@ class AwsBaseWaiterTrigger(BaseTrigger): # Should be assigned in child class, unless hook() is overridden. aws_hook_class: type[AwsGenericHook] + def __init_subclass__(cls, **kwargs: Any) -> None: + """ + Reject a subclass that cannot build a hook, at import time. + + ``hook()`` is only reached from ``run()``, which executes in the triggerer, so without this + a subclass that declares neither would defer successfully and fail later, out of sight of + the task that deferred. The operator side gets the same guarantee from + ``AwsBaseHookMixin.validate_attributes``. This runs on class creation rather than in + ``__init__`` because subclasses such as ``EksDeleteClusterTrigger`` never call + ``super().__init__()``. + """ + super().__init_subclass__(**kwargs) + if cls.hook is not AwsBaseWaiterTrigger.hook: + return + hook_class = getattr(cls, "aws_hook_class", None) + if hook_class is None: + raise AttributeError( + f"Class attribute '{cls.__name__}.aws_hook_class' should be set, " + f"or {cls.__name__}.hook() overridden." + ) + if not (isinstance(hook_class, type) and issubclass(hook_class, AwsGenericHook)): + raise AttributeError( + f"Class attribute '{cls.__name__}.aws_hook_class' is not a subclass of AwsGenericHook." + ) + def __init__( self, *, 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 index 197cdffdcd373..1f63053fa906e 100644 --- a/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py +++ b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py @@ -101,7 +101,7 @@ def find_defer_sites() -> list[tuple[str, int, str, list[str]]]: passed = {kw.arg for kw in call.keywords if kw.arg} sites.append( ( - str(path.relative_to(AWS_ROOT)), + path.relative_to(AWS_ROOT).as_posix(), call.lineno, name, [p for p in HOOK_CONFIGURATION if p not in passed], @@ -129,7 +129,7 @@ def find_unreadable_defer_sites() -> set[tuple[str, str]]: 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)) + unreadable.add((path.relative_to(AWS_ROOT).as_posix(), trigger.id)) return unreadable @@ -186,7 +186,7 @@ def find_waiter_triggers() -> list[type[AwsBaseWaiterTrigger]]: ) 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), ( + assert hasattr(trigger_class, "aws_hook_class") or trigger_class.hook is not AwsBaseWaiterTrigger.hook, ( f"{trigger_class.__name__} sets neither aws_hook_class nor hook(); " f"building its hook would fail at runtime." ) @@ -213,18 +213,18 @@ def find_hand_built_hooks() -> list[tuple[str, int, str, list[str]]]: sites = [] for path in sorted((AWS_ROOT / "triggers").rglob("*.py")): for node in ast.walk(ast.parse(path.read_text())): - if not ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id.endswith("Hook") - ): + if not isinstance(node, ast.Call): + continue + # Match a module-qualified ``module.SomeHook(...)`` as well as a bare name. + hook = node.func.attr if isinstance(node.func, ast.Attribute) else getattr(node.func, "id", "") + if not hook.endswith("Hook"): continue passed = {keyword.arg for keyword in node.keywords if keyword.arg} # AwsGenericHook names the botocore config ``config``. if "config" in passed: passed.add("botocore_config") missing = [name for name in HOOK_CONFIGURATION if name not in passed] - sites.append((path.name, node.lineno, node.func.id, missing)) + sites.append((path.name, node.lineno, hook, missing)) return sites diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py b/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py index 0c2ea69149a06..60454ee45215b 100644 --- a/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py +++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py @@ -23,10 +23,10 @@ import pytest from airflow.exceptions import AirflowException +from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger if TYPE_CHECKING: - from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook from airflow.triggers.base import TriggerEvent @@ -156,3 +156,37 @@ def test_event_from_exception(self): "message": "AWS Glue job failed.", "hello": "world", } + + +class TestAwsBaseWaiterTriggerSubclassValidation: + """``hook()`` runs in the triggerer, so a subclass that cannot build one must fail on import.""" + + def test_subclass_with_neither_hook_class_nor_hook_is_rejected(self): + with pytest.raises(AttributeError, match="aws_hook_class' should be set"): + + class MissingBoth(AwsBaseWaiterTrigger): + pass + + def test_subclass_whose_hook_class_is_not_a_hook_is_rejected(self): + with pytest.raises(AttributeError, match="not a subclass of AwsGenericHook"): + + class NotAHook(AwsBaseWaiterTrigger): + aws_hook_class = str + + def test_subclass_declaring_a_hook_class_is_accepted(self): + class Declared(AwsBaseWaiterTrigger): + aws_hook_class = AwsGenericHook + + assert Declared.aws_hook_class is AwsGenericHook + + def test_subclass_inheriting_a_hook_override_is_accepted(self): + """An override reached through an intermediate base still counts.""" + + class Intermediate(AwsBaseWaiterTrigger): + def hook(self) -> AwsGenericHook: + return AsyncMock() + + class Leaf(Intermediate): + pass + + assert Leaf.hook is Intermediate.hook From 8acadfe6acd92caf8ae14bdb129c88ca109c8010 Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Fri, 18 Sep 2026 15:33:11 +0530 Subject: [PATCH 6/9] Treat an unresolvable trigger expression as unreadable The sweep resolved a trigger= expression to the constructions it can evaluate to and returned an empty list when it could not read one. A bare reference was caught separately by matching ast.Name, but an attribute, a subscript, a conditional with one unreadable branch, or a construction whose callee cannot be named all produced nothing the checks could act on, and nothing reported their absence. The conditional was the worst of them: it still yielded the readable branch, so the site appeared in the parametrized run and read as covered. The directory filter had the same shape of problem, silently ignoring any defer site that did not sit directly in operators/ or sensors/. Reported in review by a contributor on the pull request. --- .../aws/test_deferred_hook_configuration.py | 125 +++++++++++------- 1 file changed, 78 insertions(+), 47 deletions(-) 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 index 1f63053fa906e..64b96b0c07cde 100644 --- a/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py +++ b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py @@ -21,6 +21,7 @@ import importlib import inspect import pkgutil +from collections.abc import Iterator from pathlib import Path import pytest @@ -61,21 +62,41 @@ ) -def trigger_constructions(expr: ast.expr) -> list[ast.Call]: - """Resolve a ``trigger=`` expression to the constructions it can evaluate to.""" +def read_trigger_name(call: ast.Call) -> str | None: + """The trigger class a construction names, or ``None`` when the callee cannot be read.""" + if isinstance(call.func, ast.Name): + return call.func.id + if isinstance(call.func, ast.Attribute): + return call.func.attr + return None + + +def trigger_constructions(expr: ast.expr) -> list[ast.Call] | None: + """ + Resolve a ``trigger=`` expression to the constructions it can evaluate to. + + ``None`` means the expression cannot be read statically. Returning that rather than an empty + list is what keeps a site from disappearing: a bare reference, a subscript, or a conditional + with one unreadable branch all have to be acknowledged in ``UNREADABLE_DEFER_SITES`` instead of + quietly contributing nothing to the sweep. + """ if isinstance(expr, ast.Call): - return [expr] + # A construction whose callee cannot be named is no more readable than a bare reference: + # the allowlists key on the class name, so an unnamed one could never match them. + return [expr] if read_trigger_name(expr) is not None else None if isinstance(expr, ast.IfExp): - return trigger_constructions(expr.body) + trigger_constructions(expr.orelse) - return [] + branches = (trigger_constructions(expr.body), trigger_constructions(expr.orelse)) + if any(branch is None for branch in branches): + return None + return [call for branch in branches for call in branch] + return None -def find_defer_sites() -> list[tuple[str, int, str, list[str]]]: - """Collect every ``self.defer(trigger=SomeTrigger(...))`` in the provider.""" - sites = [] +def walk_defer_sites() -> Iterator[tuple[Path, ast.expr]]: + """Yield the ``trigger=`` expression of every ``self.defer(...)`` in the provider.""" + # Every file, not just operators/ and sensors/: ``defer`` is a BaseOperator method, so a site + # can appear anywhere, and a directory filter would drop a nested subpackage without saying so. for path in sorted(AWS_ROOT.rglob("*.py")): - 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 @@ -88,49 +109,39 @@ def find_defer_sites() -> list[tuple[str, int, str, list[str]]]: ): continue trigger = next((kw.value for kw in node.keywords if kw.arg == "trigger"), None) - if trigger is None: + if trigger is not None: + yield path, trigger + + +def find_defer_sites() -> list[tuple[str, int, str, list[str]]]: + """Collect every ``self.defer(trigger=SomeTrigger(...))`` in the provider.""" + sites = [] + for path, trigger in walk_defer_sites(): + # The trigger may be built inline, or picked between in a conditional expression, so take + # every construction the expression can yield rather than assuming a single call. + for call in trigger_constructions(trigger) or (): + name = read_trigger_name(call) + if name in UNCONFIGURABLE_TRIGGERS: 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( - ( - path.relative_to(AWS_ROOT).as_posix(), - call.lineno, - name, - [p for p in HOOK_CONFIGURATION if p not in passed], - ) + passed = {kw.arg for kw in call.keywords if kw.arg} + sites.append( + ( + path.relative_to(AWS_ROOT).as_posix(), + call.lineno, + name, + [p for p in HOOK_CONFIGURATION if p not in passed], ) + ) return sites def find_unreadable_defer_sites() -> set[tuple[str, str]]: - """Defer sites whose trigger 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((path.relative_to(AWS_ROOT).as_posix(), trigger.id)) - return unreadable + """Defer sites whose trigger expression cannot be resolved to the constructions it yields.""" + return { + (path.relative_to(AWS_ROOT).as_posix(), ast.unparse(trigger)) + for path, trigger in walk_defer_sites() + if trigger_constructions(trigger) is None + } DEFER_SITES = find_defer_sites() @@ -145,6 +156,26 @@ def test_no_defer_site_escapes_the_check(): assert find_unreadable_defer_sites() == UNREADABLE_DEFER_SITES +@pytest.mark.parametrize( + ("expression", "expected"), + [ + pytest.param("SomeTrigger(x=1)", 1, id="call"), + pytest.param("A() if flag else B()", 2, id="conditional-both-readable"), + pytest.param("trigger", None, id="bare-name"), + pytest.param("self._trigger", None, id="attribute"), + pytest.param("triggers[kind]", None, id="subscript"), + pytest.param("A() if flag else self._trigger", None, id="conditional-one-unreadable"), + pytest.param("TRIGGERS[kind](x=1)", None, id="unnameable-callee"), + pytest.param("module.SomeTrigger(x=1)", 1, id="module-qualified-callee"), + ], +) +def test_unreadable_trigger_expressions_resolve_to_none(expression, expected): + """Anything the sweep cannot resolve must report None so the site is forced onto the allowlist.""" + constructions = trigger_constructions(ast.parse(expression, mode="eval").body) + + assert (constructions if constructions is None else len(constructions)) == expected + + @pytest.mark.parametrize( ("source", "line", "trigger", "missing"), DEFER_SITES, From dc419805b1626114512a7ec8fb0821f14e665d5e Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Sat, 19 Sep 2026 09:47:26 +0530 Subject: [PATCH 7/9] Resolve the trigger class name where the construction is accepted The MyPy providers job is the only red check on the branch. The sweep looked up a construction's class name separately from deciding that the construction was readable, so the name stayed optional at every use even though an unnameable callee is already rejected, and the conditional branch walk could not be narrowed. Pairing each construction with the name that made it readable removes the optionality instead of asserting it away. --- .../aws/test_deferred_hook_configuration.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) 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 index 64b96b0c07cde..04de3944c26d1 100644 --- a/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py +++ b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py @@ -71,9 +71,9 @@ def read_trigger_name(call: ast.Call) -> str | None: return None -def trigger_constructions(expr: ast.expr) -> list[ast.Call] | None: +def trigger_constructions(expr: ast.expr) -> list[tuple[ast.Call, str]] | None: """ - Resolve a ``trigger=`` expression to the constructions it can evaluate to. + Resolve a ``trigger=`` expression to the constructions it can evaluate to, each with its name. ``None`` means the expression cannot be read statically. Returning that rather than an empty list is what keeps a site from disappearing: a bare reference, a subscript, or a conditional @@ -83,12 +83,16 @@ def trigger_constructions(expr: ast.expr) -> list[ast.Call] | None: if isinstance(expr, ast.Call): # A construction whose callee cannot be named is no more readable than a bare reference: # the allowlists key on the class name, so an unnamed one could never match them. - return [expr] if read_trigger_name(expr) is not None else None + name = read_trigger_name(expr) + return [(expr, name)] if name is not None else None if isinstance(expr, ast.IfExp): - branches = (trigger_constructions(expr.body), trigger_constructions(expr.orelse)) - if any(branch is None for branch in branches): - return None - return [call for branch in branches for call in branch] + constructions: list[tuple[ast.Call, str]] = [] + for branch in (expr.body, expr.orelse): + resolved = trigger_constructions(branch) + if resolved is None: + return None + constructions.extend(resolved) + return constructions return None @@ -115,12 +119,11 @@ def walk_defer_sites() -> Iterator[tuple[Path, ast.expr]]: def find_defer_sites() -> list[tuple[str, int, str, list[str]]]: """Collect every ``self.defer(trigger=SomeTrigger(...))`` in the provider.""" - sites = [] + sites: list[tuple[str, int, str, list[str]]] = [] for path, trigger in walk_defer_sites(): # The trigger may be built inline, or picked between in a conditional expression, so take # every construction the expression can yield rather than assuming a single call. - for call in trigger_constructions(trigger) or (): - name = read_trigger_name(call) + for call, name in trigger_constructions(trigger) or (): if name in UNCONFIGURABLE_TRIGGERS: continue passed = {kw.arg for kw in call.keywords if kw.arg} From 50ce3e544058b38b8eeea27581bec0bdb0de2861 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Mon, 21 Sep 2026 16:33:55 +0200 Subject: [PATCH 8/9] Update providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py --- .../unit/amazon/aws/test_deferred_hook_configuration.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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 index 04de3944c26d1..8181274f04038 100644 --- a/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py +++ b/providers/amazon/tests/unit/amazon/aws/test_deferred_hook_configuration.py @@ -219,10 +219,11 @@ def find_waiter_triggers() -> list[type[AwsBaseWaiterTrigger]]: 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 trigger_class.hook is not AwsBaseWaiterTrigger.hook, ( - f"{trigger_class.__name__} sets neither aws_hook_class nor hook(); " - f"building its hook would fail at runtime." + """The hook a trigger names must accept what ``_hook_parameters`` will pass it.""" + if trigger_class.hook is not AwsBaseWaiterTrigger.hook: + pytest.skip(f"{trigger_class.__name__} builds its hook by hand") + inspect.signature(trigger_class.aws_hook_class).bind_partial( + aws_conn_id=None, region_name=None, verify=None, config=None ) From 32b8d5fb466ec3ed444c01f0fc93ec7bce4959cb Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Mon, 21 Sep 2026 16:38:05 +0200 Subject: [PATCH 9/9] Keep the deferred-hook warning in the pending changelog slot The 2026-09-09 provider release cut 9.36.0 and moved the Comprehend warning under that heading. Merging main carries this branch's warning down with it, attributing a change 9.36.0 does not contain to a released version. Put it back above the first version heading so it lands in the release that actually ships it. Generated-by: Claude Opus 5 --- providers/amazon/docs/changelog.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/providers/amazon/docs/changelog.rst b/providers/amazon/docs/changelog.rst index 9375efb929766..fa3f4e9ee49fe 100644 --- a/providers/amazon/docs/changelog.rst +++ b/providers/amazon/docs/changelog.rst @@ -26,9 +26,6 @@ Changelog --------- -9.36.0 -...... - .. warning:: Deferrable AWS operators and sensors now hand ``region_name``, ``verify`` and ``botocore_config`` to the trigger they defer to, so the triggerer builds its hook from the operator's settings @@ -38,6 +35,9 @@ Changelog botocore configuration, which previously never reached it. Set these explicitly on the operator if the deferred half needs to differ from the synchronous half. +9.36.0 +...... + .. warning:: The default waiter timeout of ``ComprehendCreateDocumentClassifierOperator`` was raised from 20 minutes (``waiter_max_attempts=20``) to 60 minutes (``waiter_max_attempts=60``), because