Repository: Netflix/consoleme
File: consoleme/celery/celery_tasks.py (Python)

Code before:
    )
    return True


@app.task(name="cache_roles_for_account", soft_time_limit=2700, **default_retry_kwargs)
def cache_roles_for_account(account_id: str) -> bool:
    # Get the DynamoDB handler:
    dynamo = IAMRoleDynamoHandler()
    cache_key = config.get("aws.iamroles_redis_key", "IAM_ROLE_CACHE")
    # Only query IAM and put data in Dynamo if we're in the active region
    if config.region == config.get("celery.active_region") or config.get(
        "environment"
    ) in ["dev", "test"]:
        client = boto3_cached_conn(
            "iam",
            account_number=account_id,
            assume_role=config.get("policies.role_name"),
            region=config.region,
        )
        paginator = client.get_paginator("get_account_authorization_details")
        response_iterator = paginator.paginate()
        all_iam_resources = {}
        for response in response_iterator:
            if not all_iam_resources:
                all_iam_resources = response
            all_iam_resources["UserDetailList"].extend(response["UserDetailList"])
            all_iam_resources["GroupDetailList"].extend(response["GroupDetailList"])
            all_iam_resources["RoleDetailList"].extend(response["RoleDetailList"])
            all_iam_resources["Policies"].extend(response["Policies"])
            for k in response.keys():
                if k not in [
                    "UserDetailList",
                    "GroupDetailList",
                    "RoleDetailList",
                    "Policies",
                    "ResponseMetadata",
                    "Marker",
                    "IsTruncated",
                ]:
                    # Fail hard if we find something unexpected
                    raise RuntimeError("Unexpected key {0} in response".format(k))

        # Store entire response in S3
        async_to_sync(store_json_results_in_redis_and_s3)(
            all_iam_resources,
            s3_bucket=config.get("cache_iam_resources_for_account.s3.bucket"),
            s3_key=config.get("cache_iam_resources_for_account.s3.file", "").format(
                account_id=account_id
            ),
        )


Review comment (refactor):
Should these be in an `else` block? It looks like we'll end up with duplicates for the first iteration of this `for` loop

Code after:
    )
    return True


@app.task(name="cache_roles_for_account", soft_time_limit=2700, **default_retry_kwargs)
def cache_roles_for_account(account_id: str) -> bool:
    # Get the DynamoDB handler:
    dynamo = IAMRoleDynamoHandler()
    cache_key = config.get("aws.iamroles_redis_key", "IAM_ROLE_CACHE")
    # Only query IAM and put data in Dynamo if we're in the active region
    if config.region == config.get("celery.active_region") or config.get(
        "environment"
    ) in ["dev", "test"]:
        client = boto3_cached_conn(
            "iam",
            account_number=account_id,
            assume_role=config.get("policies.role_name"),
            region=config.region,
        )
        paginator = client.get_paginator("get_account_authorization_details")
        response_iterator = paginator.paginate()
        all_iam_resources = {}
        for response in response_iterator:
            if not all_iam_resources:
                all_iam_resources = response
            else:
                all_iam_resources["UserDetailList"].extend(response["UserDetailList"])
                all_iam_resources["GroupDetailList"].extend(response["GroupDetailList"])
                all_iam_resources["RoleDetailList"].extend(response["RoleDetailList"])
                all_iam_resources["Policies"].extend(response["Policies"])
            for k in response.keys():
                if k not in [
                    "UserDetailList",
                    "GroupDetailList",
                    "RoleDetailList",
                    "Policies",
                    "ResponseMetadata",
                    "Marker",
                    "IsTruncated",
                ]:
                    # Fail hard if we find something unexpected
                    raise RuntimeError("Unexpected key {0} in response".format(k))

        # Store entire response in S3
        async_to_sync(store_json_results_in_redis_and_s3)(
            all_iam_resources,
            s3_bucket=config.get("cache_iam_resources_for_account.s3.bucket"),
            s3_key=config.get("cache_iam_resources_for_account.s3.file", "").format(
                account_id=account_id
            ),
        )