def lambda_handler()

in templates/aws-cloudfront-waf/source/log_parser/add_athena_partitions.py [0:0]
36 lines of code
10 McCabe index (conditional complexity)

def lambda_handler(event, context):
    """
    This function adds a new hourly partition to athena table.
    It runs every hour, triggered by a CloudWatch event rule.
    """
    log = logging.getLogger()
    log.debug('[add-athena-partition lambda_handler] Start')
    try:
        # ---------------------------------------------------------
        # Set Log Level
        # ---------------------------------------------------------
        log_level = str(environ['LOG_LEVEL'].upper())
        if log_level not in ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']:
            log_level = 'ERROR'
        log.setLevel(log_level)

        # ----------------------------------------------------------
        # Process event
        # ----------------------------------------------------------
        log.info(event)

        athena_client = boto3.client('athena')
        database_name = event['glueAccessLogsDatabase']
        access_log_bucket = event['accessLogBucket']
        waf_log_bucket = event['wafLogBucket']
        athena_work_group = event['athenaWorkGroup']

        try:
            # Add athena partition for cloudfront or alb logs
            if len(access_log_bucket) > 0:
                execute_athena_query(log, access_log_bucket,
                                     database_name,
                                     event['glueAppAccessLogsTable'],
                                     athena_client,
                                     athena_work_group)
        except Exception as error:
            log.error('[add-athena-partition lambda_handler] App access log Athena query execution failed: %s'%str(error))

        try:
            # Add athena partition for waf logs
            if len(waf_log_bucket) > 0:
                execute_athena_query(log, waf_log_bucket,
                                     database_name,
                                     event['glueWafAccessLogsTable'],
                                     athena_client,
                                     athena_work_group)
        except Exception as error:
            log.error('[add-athena-partition lambda_handler] WAF access log Athena query execution failed: %s'%str(error))

    except Exception as error:
        log.error(str(error))
        raise

    log.debug('[add-athena-partition lambda_handler] End')