Repository: apache/airflow
File: providers/amazon/tests/system/amazon/aws/example_emr_eks.py (Python)

Code before:

    client.update_assume_role_policy(
        RoleName=role_name,
        PolicyDocument=json.dumps(role_trust_policy),
    )

    # See https://docs.aws.amazon.com/emr/latest/EMR-on-EKS-DevelopmentGuide/setting-up-trust-policy.html
    # The action "update-role-trust-policy" is not available in boto3, thus we need to do it using AWS CLI
    commands = (
        f"aws emr-containers update-role-trust-policy --cluster-name {cluster_name} "
        f"--namespace {cluster_namespace} --role-name {role_name}"
    )

    build = subprocess.Popen(
        commands,
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    _, err = build.communicate()

    if build.returncode != 0:
        raise RuntimeError(err)

    # Wait for IAM changes to propagate to avoid authentication failures
    time.sleep(10)


@task(trigger_rule=TriggerRule.ALL_DONE)
def delete_virtual_cluster(virtual_cluster_id):
    boto3.client("emr-containers").delete_virtual_cluster(
        id=virtual_cluster_id,
    )


with DAG(
    dag_id=DAG_ID,
    schedule="@once",
    start_date=datetime(2021, 1, 1),
    catchup=False,
) as dag:
    test_context = sys_test_context_task()
    env_id = test_context[ENV_ID_KEY]
    role_arn = test_context[ROLE_ARN_KEY]
    subnets = test_context[SUBNETS_KEY]
    job_role_arn = test_context[JOB_ROLE_ARN_KEY]
    job_role_name = test_context[JOB_ROLE_NAME_KEY]

    s3_bucket_name = f"{env_id}-bucket"
    eks_cluster_name = f"{env_id}-cluster"
    virtual_cluster_name = f"{env_id}-virtual-cluster"

Review comment (suggestion):
Were you able to test if 10s is enough? We could optionally make this a variable through the context builder so that we can change the time easily without needing to PR each time.

Code after:
        if statement["Action"] != "sts:AssumeRoleWithWebIdentity"
    ]

    client.update_assume_role_policy(
        RoleName=role_name,
        PolicyDocument=json.dumps(role_trust_policy),
    )

    # See https://docs.aws.amazon.com/emr/latest/EMR-on-EKS-DevelopmentGuide/setting-up-trust-policy.html
    # The action "update-role-trust-policy" is not available in boto3, thus we need to do it using AWS CLI
    commands = (
        f"aws emr-containers update-role-trust-policy --cluster-name {cluster_name} "
        f"--namespace {cluster_namespace} --role-name {role_name}"
    )

    build = subprocess.Popen(
        commands,
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    _, err = build.communicate()

    if build.returncode != 0:
        raise RuntimeError(err)

    # Wait for IAM changes to propagate to avoid authentication failures
    time.sleep(wait_time)


@task(trigger_rule=TriggerRule.ALL_DONE)
def delete_virtual_cluster(virtual_cluster_id):
    boto3.client("emr-containers").delete_virtual_cluster(
        id=virtual_cluster_id,
    )


with DAG(
    dag_id=DAG_ID,
    schedule="@once",
    start_date=datetime(2021, 1, 1),
    catchup=False,
) as dag:
    test_context = sys_test_context_task()
    env_id = test_context[ENV_ID_KEY]
    role_arn = test_context[ROLE_ARN_KEY]
    subnets = test_context[SUBNETS_KEY]
    job_role_arn = test_context[JOB_ROLE_ARN_KEY]
    job_role_name = test_context[JOB_ROLE_NAME_KEY]
    update_trust_policy_wait_time = test_context[UPDATE_TRUST_POLICY_WAIT_TIME_KEY]