From 6157cfaa1442e6e2f0803a0c66882dbcbc5d60b6 Mon Sep 17 00:00:00 2001 From: Ovidiu Calbajos Date: Wed, 9 Aug 2023 13:05:49 +0300 Subject: [PATCH 1/3] Enhance RdsStartExportTaskOperator Add support to specify max_attempts and check_interval for RdsStartExportTaskOperator --- airflow/providers/amazon/aws/operators/rds.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/operators/rds.py b/airflow/providers/amazon/aws/operators/rds.py index 92e98d7bd64e6..12b65b731a0b1 100644 --- a/airflow/providers/amazon/aws/operators/rds.py +++ b/airflow/providers/amazon/aws/operators/rds.py @@ -351,6 +351,8 @@ def __init__( s3_prefix: str = "", export_only: list[str] | None = None, wait_for_completion: bool = True, + waiter_interval: int = 30, + waiter_max_attempts: int = 40, **kwargs, ): super().__init__(**kwargs) @@ -363,6 +365,8 @@ def __init__( self.s3_prefix = s3_prefix self.export_only = export_only or [] self.wait_for_completion = wait_for_completion + self.waiter_interval = waiter_interval + self.waiter_max_attempts = waiter_max_attempts def execute(self, context: Context) -> str: self.log.info("Starting export task %s for snapshot %s", self.export_task_identifier, self.source_arn) @@ -378,7 +382,12 @@ def execute(self, context: Context) -> str: ) if self.wait_for_completion: - self.hook.wait_for_export_task_state(self.export_task_identifier, target_state="complete") + self.hook.wait_for_export_task_state( + export_task_id=self.export_task_identifier, + target_state="complete", + check_interval=self.waiter_interval, + max_attempts=self.waiter_max_attempts + ) return json.dumps(start_export, default=str) From ac8575af5656b128264993f0e6d5040262d045c0 Mon Sep 17 00:00:00 2001 From: Ovidiu Calbajos Date: Wed, 9 Aug 2023 14:41:36 +0300 Subject: [PATCH 2/3] Fix typo --- airflow/providers/amazon/aws/operators/rds.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/operators/rds.py b/airflow/providers/amazon/aws/operators/rds.py index 12b65b731a0b1..0951eb36ea476 100644 --- a/airflow/providers/amazon/aws/operators/rds.py +++ b/airflow/providers/amazon/aws/operators/rds.py @@ -386,7 +386,7 @@ def execute(self, context: Context) -> str: export_task_id=self.export_task_identifier, target_state="complete", check_interval=self.waiter_interval, - max_attempts=self.waiter_max_attempts + max_attempts=self.waiter_max_attempts, ) return json.dumps(start_export, default=str) From 1f9263f3994a28420f482ef16fbc4e9ccd6a335d Mon Sep 17 00:00:00 2001 From: Ovidiu Calbajos Date: Wed, 9 Aug 2023 22:32:46 +0300 Subject: [PATCH 3/3] Update docstring --- airflow/providers/amazon/aws/operators/rds.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/airflow/providers/amazon/aws/operators/rds.py b/airflow/providers/amazon/aws/operators/rds.py index 0951eb36ea476..f7361892b059c 100644 --- a/airflow/providers/amazon/aws/operators/rds.py +++ b/airflow/providers/amazon/aws/operators/rds.py @@ -328,6 +328,8 @@ class RdsStartExportTaskOperator(RdsBaseOperator): :param s3_prefix: The Amazon S3 bucket prefix to use as the file name and path of the exported snapshot. :param export_only: The data to be exported from the snapshot. :param wait_for_completion: If True, waits for the DB snapshot export to complete. (default: True) + :param waiter_interval: The number of seconds to wait before checking the export status. (default: 30) + :param waiter_max_attempts: The number of attempts to make before failing. (default: 40) """ template_fields = (