Skip to content

Fix RdsStopDbOperator operator in deferrable mode #41059

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vincbeck merged 2 commits into from
Jul 26, 2024
Merged

Fix RdsStopDbOperator operator in deferrable mode #41059

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
da94f6f
Fix `RdsStopDbOperator` operator in deferrable mode
vincbeck Jul 26, 2024
9f92028
Fix tests
vincbeck Jul 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions airflow/providers/amazon/aws/hooks/rds.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,7 +259,7 @@ def poke():
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,
Expand All @@ -272,7 +272,7 @@ def poke():
)
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:
"""
Expand Down Expand Up @@ -310,7 +310,7 @@ def poke():
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,
Expand Down
37 changes: 17 additions & 20 deletions airflow/providers/amazon/aws/operators/rds.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
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
Expand Down Expand Up @@ -782,7 +783,7 @@ def execute(self, context: Context) -> str:
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",
)
Expand Down Expand Up @@ -881,12 +882,25 @@ def execute(self, context: Context) -> str:
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:
Expand Down Expand Up @@ -915,23 +929,6 @@ def _stop_db(self):
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",
Expand Down
253 changes: 253 additions & 0 deletions airflow/providers/amazon/aws/waiters/rds.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
}
]
}
}
}
Loading