This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 83ca61a501 Fix `RdsStopDbOperator` operator in deferrable mode (#41059)
83ca61a501 is described below
commit 83ca61a501d755669fc83b1ad9038d0ca9d600ad
Author: Vincent <97131062+vincb...@users.noreply.github.com>
AuthorDate: Fri Jul 26 16:52:25 2024 -0400
Fix `RdsStopDbOperator` operator in deferrable mode (#41059)
---
airflow/providers/amazon/aws/hooks/rds.py | 6 +-
airflow/providers/amazon/aws/operators/rds.py | 37 ++--
airflow/providers/amazon/aws/waiters/rds.json | 253 +++++++++++++++++++++++
tests/providers/amazon/aws/hooks/test_rds.py | 24 +--
tests/providers/amazon/aws/operators/test_rds.py | 14 +-
5 files changed, 280 insertions(+), 54 deletions(-)
diff --git a/airflow/providers/amazon/aws/hooks/rds.py
b/airflow/providers/amazon/aws/hooks/rds.py
index 8219a37757..588d78c782 100644
--- a/airflow/providers/amazon/aws/hooks/rds.py
+++ b/airflow/providers/amazon/aws/hooks/rds.py
@@ -259,7 +259,7 @@ class RdsHook(AwsGenericHook["RDSClient"]):
return self.get_db_instance_state(db_instance_id)
target_state = target_state.lower()
- if target_state in ("available", "deleted"):
+ if target_state in ("available", "deleted", "stopped"):
waiter = self.conn.get_waiter(f"db_instance_{target_state}") #
type: ignore
wait(
waiter=waiter,
@@ -272,7 +272,7 @@ class RdsHook(AwsGenericHook["RDSClient"]):
)
else:
self._wait_for_state(poke, target_state, check_interval,
max_attempts)
- self.log.info("DB cluster snapshot '%s' reached the '%s' state",
db_instance_id, target_state)
+ self.log.info("DB cluster '%s' reached the '%s' state",
db_instance_id, target_state)
def get_db_cluster_state(self, db_cluster_id: str) -> str:
"""
@@ -310,7 +310,7 @@ class RdsHook(AwsGenericHook["RDSClient"]):
return self.get_db_cluster_state(db_cluster_id)
target_state = target_state.lower()
- if target_state in ("available", "deleted"):
+ if target_state in ("available", "deleted", "stopped"):
waiter = self.conn.get_waiter(f"db_cluster_{target_state}") #
type: ignore
waiter.wait(
DBClusterIdentifier=db_cluster_id,
diff --git a/airflow/providers/amazon/aws/operators/rds.py
b/airflow/providers/amazon/aws/operators/rds.py
index a2f35b5081..f37c698d87 100644
--- a/airflow/providers/amazon/aws/operators/rds.py
+++ b/airflow/providers/amazon/aws/operators/rds.py
@@ -36,6 +36,7 @@ from airflow.providers.amazon.aws.utils import
validate_execute_complete_event
from airflow.providers.amazon.aws.utils.rds import RdsDbType
from airflow.providers.amazon.aws.utils.tags import format_tags
from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
+from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
from mypy_boto3_rds.type_defs import TagTypeDef
@@ -782,7 +783,7 @@ class RdsStartDbOperator(RdsBaseOperator):
aws_conn_id=self.aws_conn_id,
region_name=self.region_name,
response=start_db_response,
- db_type=RdsDbType.INSTANCE,
+ db_type=self.db_type,
),
method_name="execute_complete",
)
@@ -881,12 +882,25 @@ class RdsStopDbOperator(RdsBaseOperator):
aws_conn_id=self.aws_conn_id,
region_name=self.region_name,
response=stop_db_response,
- db_type=RdsDbType.INSTANCE,
+ db_type=self.db_type,
),
method_name="execute_complete",
)
elif self.wait_for_completion:
- self._wait_until_db_stopped()
+ waiter = self.hook.get_waiter(f"db_{self.db_type.value}_stopped")
+ waiter_key = (
+ "DBInstanceIdentifier" if self.db_type == RdsDbType.INSTANCE
else "DBClusterIdentifier"
+ )
+ kwargs = {waiter_key: self.db_identifier}
+ waiter.wait(
+ WaiterConfig=prune_dict(
+ {
+ "Delay": self.waiter_delay,
+ "MaxAttempts": self.waiter_max_attempts,
+ }
+ ),
+ **kwargs,
+ )
return json.dumps(stop_db_response, default=str)
def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> str:
@@ -915,23 +929,6 @@ class RdsStopDbOperator(RdsBaseOperator):
response =
self.hook.conn.stop_db_cluster(DBClusterIdentifier=self.db_identifier)
return response
- def _wait_until_db_stopped(self):
- self.log.info("Waiting for DB %s to reach 'stopped' state",
self.db_type.value)
- if self.db_type == RdsDbType.INSTANCE:
- self.hook.wait_for_db_instance_state(
- self.db_identifier,
- target_state="stopped",
- check_interval=self.waiter_delay,
- max_attempts=self.waiter_max_attempts,
- )
- else:
- self.hook.wait_for_db_cluster_state(
- self.db_identifier,
- target_state="stopped",
- check_interval=self.waiter_delay,
- max_attempts=self.waiter_max_attempts,
- )
-
__all__ = [
"RdsCreateDbSnapshotOperator",
diff --git a/airflow/providers/amazon/aws/waiters/rds.json
b/airflow/providers/amazon/aws/waiters/rds.json
new file mode 100644
index 0000000000..78c56da53f
--- /dev/null
+++ b/airflow/providers/amazon/aws/waiters/rds.json
@@ -0,0 +1,253 @@
+{
+ "version": 2,
+ "waiters": {
+ "db_instance_stopped": {
+ "operation": "DescribeDBInstances",
+ "delay": 30,
+ "maxAttempts": 60,
+ "acceptors": [
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "stopped",
+ "state": "success"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "available",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "backing-up",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "creating",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "delete-precheck",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "deleting",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "failed",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "inaccessible-encryption-credentials",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected":
"inaccessible-encryption-credentials-recoverable",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "incompatible-network",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "incompatible-option-group",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "incompatible-parameters",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "incompatible-restore",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "insufficient-capacity",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "maintenance",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "modifying",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "rebooting",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "renaming",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "restore-error",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "starting",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "stopping",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "storage-full",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "upgrading",
+ "state": "retry"
+ }
+ ]
+ },
+ "db_cluster_stopped": {
+ "operation": "DescribeDBClusters",
+ "delay": 30,
+ "maxAttempts": 60,
+ "acceptors": [
+ {
+ "matcher": "pathAll",
+ "argument": "DBClusters[].Status",
+ "expected": "stopped",
+ "state": "success"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "available",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "backing-up",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "cloning-failed",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "creating",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "deleting",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "failing-over",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "inaccessible-encryption-credentials",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected":
"inaccessible-encryption-credentials-recoverable",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "maintenance",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "migrating",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "migration-failed",
+ "state": "failure"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "modifying",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "renaming",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "starting",
+ "state": "retry"
+ },
+ {
+ "matcher": "pathAll",
+ "argument": "DBInstances[].DBInstanceStatus",
+ "expected": "stopping",
+ "state": "retry"
+ }
+ ]
+ }
+ }
+}
diff --git a/tests/providers/amazon/aws/hooks/test_rds.py
b/tests/providers/amazon/aws/hooks/test_rds.py
index b2668febfb..77159b4e14 100644
--- a/tests/providers/amazon/aws/hooks/test_rds.py
+++ b/tests/providers/amazon/aws/hooks/test_rds.py
@@ -150,7 +150,7 @@ class TestRdsHook:
def test_wait_for_db_instance_state_boto_waiters(self, rds_hook: RdsHook,
db_instance_id: str):
"""Checks that the DB instance waiter uses AWS boto waiters where
possible"""
- for state in ("available", "deleted"):
+ for state in ("available", "deleted", "stopped"):
with patch.object(rds_hook.conn, "get_waiter") as mock:
rds_hook.wait_for_db_instance_state(db_instance_id,
target_state=state, **self.waiter_args)
mock.assert_called_once_with(f"db_instance_{state}")
@@ -161,16 +161,6 @@ class TestRdsHook:
},
)
- def test_wait_for_db_instance_state_custom_waiter(self, rds_hook: RdsHook,
db_instance_id: str):
- """Checks that the DB instance waiter uses custom wait logic when AWS
boto waiters aren't available"""
- with patch.object(rds_hook, "_wait_for_state") as mock:
- rds_hook.wait_for_db_instance_state(db_instance_id,
target_state="stopped", **self.waiter_args)
- mock.assert_called_once()
-
- with patch.object(rds_hook, "get_db_instance_state",
return_value="stopped") as mock:
- rds_hook.wait_for_db_instance_state(db_instance_id,
target_state="stopped", **self.waiter_args)
- mock.assert_called_once_with(db_instance_id)
-
def test_get_db_cluster_state(self, rds_hook: RdsHook, db_cluster_id: str):
response =
rds_hook.conn.describe_db_clusters(DBClusterIdentifier=db_cluster_id)
state_expected = response["DBClusters"][0]["Status"]
@@ -179,7 +169,7 @@ class TestRdsHook:
def test_wait_for_db_cluster_state_boto_waiters(self, rds_hook: RdsHook,
db_cluster_id: str):
"""Checks that the DB cluster waiter uses AWS boto waiters where
possible"""
- for state in ("available", "deleted"):
+ for state in ("available", "deleted", "stopped"):
with patch.object(rds_hook.conn, "get_waiter") as mock:
rds_hook.wait_for_db_cluster_state(db_cluster_id,
target_state=state, **self.waiter_args)
mock.assert_called_once_with(f"db_cluster_{state}")
@@ -191,16 +181,6 @@ class TestRdsHook:
},
)
- def test_wait_for_db_cluster_state_custom_waiter(self, rds_hook: RdsHook,
db_cluster_id: str):
- """Checks that the DB cluster waiter uses custom wait logic when AWS
boto waiters aren't available"""
- with patch.object(rds_hook, "_wait_for_state") as mock_wait_for_state:
- rds_hook.wait_for_db_cluster_state(db_cluster_id,
target_state="stopped", **self.waiter_args)
- mock_wait_for_state.assert_called_once()
-
- with patch.object(rds_hook, "get_db_cluster_state",
return_value="stopped") as mock:
- rds_hook.wait_for_db_cluster_state(db_cluster_id,
target_state="stopped", **self.waiter_args)
- mock.assert_called_once_with(db_cluster_id)
-
def test_get_db_snapshot_state(self, rds_hook: RdsHook, db_snapshot_id:
str):
response =
rds_hook.conn.describe_db_snapshots(DBSnapshotIdentifier=db_snapshot_id)
state_expected = response["DBSnapshots"][0]["Status"]
diff --git a/tests/providers/amazon/aws/operators/test_rds.py
b/tests/providers/amazon/aws/operators/test_rds.py
index fd464019dd..651db53d42 100644
--- a/tests/providers/amazon/aws/operators/test_rds.py
+++ b/tests/providers/amazon/aws/operators/test_rds.py
@@ -813,8 +813,7 @@ class TestRdsStopDbOperator:
del cls.hook
@mock_aws
- @patch.object(RdsHook, "wait_for_db_instance_state")
- def test_stop_db_instance(self, mock_await_status):
+ def test_stop_db_instance(self):
_create_db_instance(self.hook)
stop_db_instance = RdsStopDbOperator(task_id="test_stop_db_instance",
db_identifier=DB_INSTANCE_NAME)
_patch_hook_get_connection(stop_db_instance.hook)
@@ -822,11 +821,10 @@ class TestRdsStopDbOperator:
result =
self.hook.conn.describe_db_instances(DBInstanceIdentifier=DB_INSTANCE_NAME)
status = result["DBInstances"][0]["DBInstanceStatus"]
assert status == "stopped"
- mock_await_status.assert_called()
@mock_aws
- @patch.object(RdsHook, "wait_for_db_instance_state")
- def test_stop_db_instance_no_wait(self, mock_await_status):
+ @patch.object(RdsHook, "get_waiter")
+ def test_stop_db_instance_no_wait(self, mock_get_waiter):
_create_db_instance(self.hook)
stop_db_instance = RdsStopDbOperator(
task_id="test_stop_db_instance_no_wait",
db_identifier=DB_INSTANCE_NAME, wait_for_completion=False
@@ -836,7 +834,7 @@ class TestRdsStopDbOperator:
result =
self.hook.conn.describe_db_instances(DBInstanceIdentifier=DB_INSTANCE_NAME)
status = result["DBInstances"][0]["DBInstanceStatus"]
assert status == "stopped"
- mock_await_status.assert_not_called()
+ mock_get_waiter.assert_not_called()
@mock.patch.object(RdsHook, "conn")
def test_deferred(self, conn_mock):
@@ -872,8 +870,7 @@ class TestRdsStopDbOperator:
assert len(instance_snapshots) == 1
@mock_aws
- @patch.object(RdsHook, "wait_for_db_cluster_state")
- def test_stop_db_cluster(self, mock_await_status):
+ def test_stop_db_cluster(self):
_create_db_cluster(self.hook)
stop_db_cluster = RdsStopDbOperator(
task_id="test_stop_db_cluster", db_identifier=DB_CLUSTER_NAME,
db_type="cluster"
@@ -884,7 +881,6 @@ class TestRdsStopDbOperator:
describe_result =
self.hook.conn.describe_db_clusters(DBClusterIdentifier=DB_CLUSTER_NAME)
status = describe_result["DBClusters"][0]["Status"]
assert status == "stopped"
- mock_await_status.assert_called()
@mock_aws
def test_stop_db_cluster_create_snapshot_logs_warning_message(self,
caplog):
