Merge pull request #240 from awslabs/more-robust-data-collection

marke cases and budgets collection more robust

Author: esc1144

data-collection/deploy/account-collector.yaml

@@ -146,7 +146,7 @@ Resources:
                     raise Exception(f"Lambda event must have 'Type' parameter with value = ({list(functions.keys())})") #pylint: disable=broad-exception-raised
 
                 account_iterator = functions[account_type]
-                with open(TMP_FILE, "w") as f:
+                with open(TMP_FILE, "w", encoding='utf-8') as f:
                     count = 0
                     f.write("[\n")
                     for account in account_iterator():
@@ -194,10 +194,9 @@ Resources:
                                 account = account_data.split(',')
                                 yield format_account(account[0], account[1], account[2])
                     else:
-                        logger.info(f'Using payer organization for the account list')
+                        logger.info('Using payer organization for the account list')
                         excluded_accounts = get_from_bucket(BUCKET, EXCLUDED_ACCOUNT_LIST_KEY)
                         if excluded_accounts:
-                            pass
                             logger.info(f'Found list of accounts to exclude in s3://{BUCKET}/{EXCLUDED_ACCOUNT_LIST_KEY}. Will only collect accounts that are not in the list')
                             excluded_accounts = [a.strip() for a in excluded_accounts[0].split(',') if a]
                         for org_account_data in iterate_admins_accounts('organizations'):
@@ -228,7 +227,7 @@ Resources:
                 try:
                     data = s3.get_object(Bucket=bucket, Key=key)
                     return data['Body'].read().decode('utf-8').strip('\n').split('\n')
-                except Exception as exc: #pylint: disable=broad-exception-caught
+                except Exception: #pylint: disable=broad-exception-caught
                     return None
 
             def format_account(account_id, account_name, payer_id):

data-collection/deploy/module-budgets.yaml

@@ -110,7 +110,6 @@ Resources:
           import logging
           import datetime
           from json import JSONEncoder
-          import sys
           
           import boto3
 
@@ -157,31 +156,36 @@ Resources:
               payer_id = account["payer_id"]
 
               logger.info(f"Collecting data for account: {account_id}")
-              budgets_client = assume_role(account_id, "budgets", "us-east-1") # must be us-east-1
-              count = 0
-              with open(TMP_FILE, "w", encoding='utf-8') as f:
-                  for budget in budgets_client.get_paginator("describe_budgets").paginate(AccountId=account_id).search('Budgets'):
-                      if not budget: continue
-                      budget['collection_time'] = collection_time
-
-                      # Fetch tags for the budget using List tag for resource API
-                      budget_name = budget['BudgetName']
-                      budget_tags = budgets_client.list_tags_for_resource(ResourceARN=f"arn:{aws_partition}:budgets::{account_id}:budget/{budget_name}")
-                      budget.update({
-                          'Account_ID': account_id,
-                          'Account_Name': account_name,
-                          'Tags': budget_tags.get('ResourceTags') or []
-                      })
-
-                      # Fetch CostFilters if available
-                      if not budget.get('CostFilters'):
-                          budget.update({'CostFilters': {'Filter': ['None']}})
-
-                      f.write(json.dumps(budget, cls=DateTimeEncoder) + "\n")
-                      count += 1
-              logger.info(f"Budgets collected: {count}")
-              s3_upload(account_id, payer_id)
+              try:
+                  budgets_client = assume_role(account_id, "budgets", "us-east-1") # must be us-east-1
+                  count = 0
+                  with open(TMP_FILE, "w", encoding='utf-8') as f:
+                      for budget in budgets_client.get_paginator("describe_budgets").paginate(AccountId=account_id).search('Budgets'):
+                          if not budget:
+                              continue
+                          budget['collection_time'] = collection_time
+                          # Fetch tags for the budget using List tag for resource API
+                          budget_name = budget['BudgetName']
+                          budget_tags = budgets_client.list_tags_for_resource(ResourceARN=f"arn:{aws_partition}:budgets::{account_id}:budget/{budget_name}")
+                          budget.update({
+                              'Account_ID': account_id,
+                              'Account_Name': account_name,
+                              'Tags': budget_tags.get('ResourceTags') or []
+                          })
+                          # Fetch CostFilters if available
+                          if not budget.get('CostFilters'):
+                              budget.update({'CostFilters': {'Filter': ['None']}})
 
+                          f.write(json.dumps(budget, cls=DateTimeEncoder) + "\n")
+                          count += 1
+                  logger.info(f"Budgets collected: {count}")
+                  s3_upload(account_id, payer_id)
+              except Exception as exc: #pylint: disable=broad-exception-caught
+                  if "AccessDenied" in str(exc):
+                      print(f'Failed to assume role {ROLE_NAME} in account {account_id}. Please make sure the role exists. {exc}')
+                  else:
+                      print(f'{exc}. Gracefully exiting from Lambda so we do not break all StepFunction Execution')
+                  return
 
           def s3_upload(account_id, payer_id):
               if os.path.getsize(TMP_FILE) == 0:

data-collection/deploy/module-support-cases.yaml

@@ -143,7 +143,17 @@ Resources:
                   raise RuntimeError(f"(MissingParameterError) Lambda event missing '{key}' parameter")
 
               account = json.loads(event[key])
-              main(account, ROLE_NAME, MODULE_NAME, BUCKET)
+
+              try:
+                  main(account, ROLE_NAME, MODULE_NAME, BUCKET)
+              except Exception as exc: #pylint: disable=broad-exception-caught
+                  if "AccessDenied" in str(exc):
+                      print(f'Failed to assume role {ROLE_NAME} in account {account}. Please make sure the role exists. {exc}')
+                  else:
+                      print(f'{exc}. Gracefully exiting from Lambda so we do not break all StepFunction Execution')
+                  return {
+                      'statusCode': 200
+                  }
               return {
                   'statusCode': 200
               }
@@ -172,7 +182,7 @@ Resources:
                       x.isoformat() if isinstance(x, (date, datetime)) else None
               )
 
-          def main(account, role_name, module_name, bucket):
+          def main(account, role_name, module_name, bucket): #pylint: disable=too-many-locals
               account_id = account["account_id"]
               payer_id = account["payer_id"]
               account_name = account.get("account_name", None)

data-collection/test/test_from_scratch.py

@@ -41,6 +41,8 @@
 logger = logging.getLogger(__name__)
 
 
+def test_deployment_works(athena):
+    pass
 
 def test_budgets_data(athena):
     data = athena_query(athena=athena, sql_query='SELECT * FROM "optimization_data"."budgets_data" LIMIT 10;')