diff --git a/astronomer/providers/amazon/aws/example_dags/example_aws_nuke.py b/astronomer/providers/amazon/aws/example_dags/example_aws_nuke.py index e41185baa..eb6950f54 100644 --- a/astronomer/providers/amazon/aws/example_dags/example_aws_nuke.py +++ b/astronomer/providers/amazon/aws/example_dags/example_aws_nuke.py @@ -22,6 +22,12 @@ SLACK_USERNAME = os.getenv("SLACK_USERNAME", "airflow_app") SLACK_WEBHOOK_CONN = os.getenv("SLACK_WEBHOOK_CONN", "http_slack") +REGRESSION_CLUSTER_AWS_ACCESS_KEY = os.getenv("REGRESSION_CLUSTER_AWS_ACCESS_KEY", "**********") +REGRESSION_CLUSTER_AWS_SECRET_ACCESS_KEY = os.getenv( + "REGRESSION_CLUSTER_AWS_SECRET_ACCESS_KEY", "***********" +) +REGRESSION_CLUSTER_AWS_DEFAULT_REGION = os.getenv("REGRESSION_CLUSTER_AWS_DEFAULT_REGION", "us-east-1") + def generate_task_report(**context: Any) -> None: """Generate a report of the task statuses for the DAG run and send it to configured Slack channel for alerts.""" @@ -125,6 +131,15 @@ def check_dag_status(**kwargs: Any) -> None: f"aws emr-containers list-virtual-clusters --state RUNNING --region {AWS_DEFAULT_REGION} | jq -r '.virtualClusters[].id' | xargs -I % aws emr-containers delete-virtual-cluster --id % --region {AWS_DEFAULT_REGION}; ", ) + terminate_dag_authoring_regression_clusters = BashOperator( + task_id="terminate_dag_authoring_regression_clusters", + bash_command=f"set -e; " + f"aws configure set aws_access_key_id {REGRESSION_CLUSTER_AWS_ACCESS_KEY}; " + f"aws configure set aws_secret_access_key {REGRESSION_CLUSTER_AWS_SECRET_ACCESS_KEY}; " + f"aws configure set default.region {REGRESSION_CLUSTER_AWS_DEFAULT_REGION}; " + f"sh $AIRFLOW_HOME/dags/example_delete_eks_cluster_and_nodes.sh {REGRESSION_CLUSTER_AWS_DEFAULT_REGION}", + ) + execute_aws_nuke = BashOperator( task_id="execute_aws_nuke", bash_command=f"aws configure set aws_access_key_id {AWS_ACCESS_KEY_ID}; " @@ -163,6 +178,7 @@ def check_dag_status(**kwargs: Any) -> None: start >> [get_airflow_version, get_airflow_executor] >> terminate_running_emr_virtual_clusters + >> terminate_dag_authoring_regression_clusters >> execute_aws_nuke >> delete_stale_emr_vpcs >> delete_stale_emr_iam_roles diff --git a/astronomer/providers/amazon/aws/example_dags/example_delete_eks_cluster_and_nodes.sh b/astronomer/providers/amazon/aws/example_dags/example_delete_eks_cluster_and_nodes.sh new file mode 100755 index 000000000..1e58da0f5 --- /dev/null +++ b/astronomer/providers/amazon/aws/example_dags/example_delete_eks_cluster_and_nodes.sh @@ -0,0 +1,39 @@ +#!/bin/bash + +# Check if the region parameter is provided +if [ -z "$1" ]; then + echo "Usage: $0 " + exit 1 +fi + +AWS_REGION="$1" + +# List all EKS clusters +clusters=$(aws eks list-clusters --region $AWS_REGION | jq -r '.clusters[]') + +# Loop through each EKS cluster +for cluster in $clusters; do + echo "Processing EKS cluster: $cluster" + + # List nodegroups for the cluster + nodegroups=$(aws eks list-nodegroups --cluster-name $cluster --region $AWS_REGION | jq -r '.nodegroups[]') + + # Delete each nodegroup + for nodegroup in $nodegroups; do + echo "Deleting nodegroup '$nodegroup' for cluster '$cluster'" + aws eks delete-nodegroup --cluster-name $cluster --nodegroup-name $nodegroup --region $AWS_REGION + aws eks wait nodegroup-deleted --cluster-name $cluster --nodegroup-name $nodegroup --region $AWS_REGION + done + + + # Delete the EKS cluster + echo "Deleting EKS cluster: $cluster" + aws eks delete-cluster --name $cluster --region $AWS_REGION + + # Wait for the EKS cluster to be deleted + echo "Waiting for EKS cluster '$cluster' to be deleted..." + aws eks wait cluster-deleted --name $cluster --region $AWS_REGION + +done + +echo "Script execution completed."