将CSV格式的AWS CUR处理后导入到AWS Athena

# 此脚本主要做两件事
# 1. 解压并读取原始S3路径中的.csv.gz格式的CUR文件，将字段名中的"/"修改为"_"，将处理完成后的CUR文件以.csv格式存储到目标S3路径
# 2. 创建Glue Crawler将处理完成的CUR文件导入到Athena中

import boto3
import pandas as pd

# 初始化变量
src_s3_bucket = "src_bucketname"  # 原始CUR存储桶名称
src_s3_prefix = "prefix"  # 原始CUR存储桶前缀
des_s3_path = "s3://des_bucketname/curtest/"  # 目标CUR存储桶路径，处理完的CUR保存路径

glue_crawler_name = "crawler-test"   # 自定义glue crawler的名称
glue_crawler_role_arn = "arn:aws:iam::xxxxxxxxxxxx:role/service-role/AWSGlueServiceRole-test"  # 设置glue crawler需要使用的IAM Role
glue_database_name = "gluedb-test"  # 填写glue database的名称，

# 初始化S3对象
s3 = boto3.client('s3')

# 读取s3中的cur对象清单
s3_objects = s3.list_objects(
    Bucket = src_s3_bucket,
    Prefix = src_s3_prefix
)

# 读取cur对象名，并写入一个列表
contents_counts = len(s3_objects['Contents'])
cur_list = []
for i in range(contents_counts):
    # 用Size!=0排除list_objects输出的s3纯路径，只保留对象，并筛选后缀是csv.gz的文件，没有想到其他好办法，暂时用这个方法吧
    if s3_objects['Contents'][i]['Size'] != 0 and s3_objects['Contents'][i]['Key'].find("csv.gz") > 0:
        cur_list.append(s3_objects['Contents'][i]['Key'])

# 将所有原始cur文件中的"/"替换为"_"（athena查询命令不支持”/"），并将处理后的文件写入目的s3
cur_counts = len(cur_list)
for i in range(cur_counts):
    cur_s3_uri = "s3://" + src_s3_bucket + "/" + cur_list[i]
    # 读取s3中的CUR文件
    cur_df = pd.read_csv(cur_s3_uri)
    # 修改CUR字段名中的 “/” 修改为 “_”
    cur_df.columns = cur_df.columns.str.replace(r"/", "_")
    # 读取CUR文件名
    cur_file_name = cur_s3_uri[(cur_s3_uri.rfind("/") + 1):-3]
    # 将处理后的CUR文件写入S3存储桶，对象名称中加一个str(i)是为了防止对象名称重叠
    cur_df.to_csv(des_s3_path + str(i) + cur_file_name, encoding='utf-8-sig', index=False)

# 使用Amazon Glue将处理后的cur导入Amazon Athena
# 初始化Glue对象
glue = boto3.client('glue')

# 如果Glue Database不存在，则新创建一个Database，没想到更好的判断数据库是否存在的方法，暂时用这个吧
try:
    glue.get_database(Name=glue_database_name)
    glue_database_exist = True
except:
    glue_database_exist = False
if not glue_database_exist:
    create_database_response = glue.create_database(
        DatabaseInput = {
            "Name": glue_database_name
        }
    )

# 创建Glue Crawler
create_crawler_response = glue.create_crawler(
                            Name = glue_crawler_name,
                            Role = glue_crawler_role_arn,
                            DatabaseName = glue_database_name,
                            Targets = {
                                'S3Targets': [
                                    {
                                        'Path': des_s3_path
                                    }
                                ]
                            }
                        )

# 启动Glue Crawler
start_crawler =  glue.start_crawler(
    Name = glue_crawler_name
)

#作者：Zhao Yipeng