增强Amazon Athena对历史查询记录的统计分析功能

import boto3
import time
import io
import json
import uuid
from datetime import datetime, date
def lambda_handler(event, context):
    # TODO implement
    to_scrape=5000
    page_size=50
    bucket = "quandata1"
    prefix = "athena_his/"
    s3_client = boto3.client('s3')
    ath = boto3.client('athena')
    paginator = ath.get_paginator("list_query_executions")
    dynamodb = boto3.resource('dynamodb')
    def json_serial(obj):
        if isinstance(obj, (datetime, date)):
            return obj.isoformat()
        raise TypeError("Type %s not serializable" % type(obj))
    df = []
    break_flag = False
    for workgroup in [w['Name'] for w in ath.list_work_groups()['WorkGroups']]:
        print(f'running {workgroup}')
        i=0
        table = dynamodb.Table('athena_his_sql')
        response = table.get_item(Key={'workgroup': workgroup,})
        curr_query_id = ''
        if  ("Item" in response): curr_query_id = response['Item']['curr_query_id']
        print("get ddb curr_id: " + workgroup + curr_query_id)
        args = {"WorkGroup": workgroup, "PaginationConfig": {"MaxItems": to_scrape, "PageSize": page_size}}
        for page in paginator.paginate(**args):
            query_ids = page['QueryExecutionIds']
            for query_id in query_ids:
                print("query_id:" + query_id)
                if i == 0:
                    table.update_item(
                        Key={'workgroup': workgroup,},
                        UpdateExpression='SET curr_query_id = :val1',
                        ExpressionAttributeValues={':val1': query_id})
                    print("update ddb curr_id: " + query_id)
                if query_id == curr_query_id:
                    break_flag=True
                    break
                query_metadata = ath.get_query_execution(QueryExecutionId=query_id)['QueryExecution']
                df.append(query_metadata)
                i += 1
            if break_flag==True:
                break
            time.sleep(1)
    json_writer = io.BytesIO()
    for record in df:
        line = json.dumps(record, default=json_serial) + "\n"
        json_writer.write(line.encode())
    s3_client.put_object(
        Bucket=bucket,
        Key=prefix + "%s.json" % uuid.uuid4(),
        ContentType='text/json',
        Body=json_writer.getvalue()
)

CREATE EXTERNAL TABLE athena_queries (
    QueryExecutionId string,
    Query string,
    StatementType string,
    Status struct<State:string,SubmissionDateTime:string,CompletionDateTime:string>,
    Statistics struct<EngineExecutionTimeInMillis:int,DataScannedInBytes:int, TotalExecutionTimeInMillis:int, QueryQueueTimeInMillis:int, QueryPlanningTimeInMillis:int, ServiceProcessingTimeInMillis:int>,
    WorkGroup string
)
ROW FORMAT SERDE 'org.openx.data.jsonserde.JsonSerDe'
LOCATION 's3://quandata1/athena_his';

CREATE EXTERNAL TABLE cloudtrail_logs(
    eventVersion STRING,
    userIdentity STRUCT<
        type: STRING,
        principalId: STRING,
        arn: STRING,
        accountId: STRING,
        invokedBy: STRING,
        accessKeyId: STRING,
        userName: STRING,
        sessionContext: STRUCT<
            attributes: STRUCT<
                mfaAuthenticated: STRING,
                creationDate: STRING>,
            sessionIssuer: STRUCT<
                type: STRING,
                principalId: STRING,
                arn: STRING,
                accountId: STRING,
                userName: STRING>>>,
    eventTime STRING,
    eventSource STRING,
    eventName STRING,
    awsRegion STRING,
    sourceIpAddress STRING,
    userAgent STRING,
    errorCode STRING,
    errorMessage STRING,
    requestParameters STRING,
    responseElements STRING,
    additionalEventData STRING,
    requestId STRING,
    eventId STRING,
    readOnly STRING,
    resources ARRAY<STRUCT<
        arn: STRING,
        accountId: STRING,
        type: STRING>>,
    eventType STRING,
    apiVersion STRING,
    recipientAccountId STRING,
    serviceEventDetails STRING,
    sharedEventID STRING,
    vpcEndpointId STRING
  )
PARTITIONED BY (
   `timestamp` string)
ROW FORMAT SERDE 'com.amazon.emr.hive.serde.CloudTrailSerde'
STORED AS INPUTFORMAT 'com.amazon.emr.cloudtrail.CloudTrailInputFormat'
OUTPUTFORMAT 'org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat'
LOCATION
  's3://bucket/AWSLogs/account-id/CloudTrail/aws-region'
TBLPROPERTIES (
  'projection.enabled'='true', 
  'projection.timestamp.format'='yyyy/MM/dd', 
  'projection.timestamp.interval'='1', 
  'projection.timestamp.interval.unit'='DAYS', 
  'projection.timestamp.range'='2021/01/01,NOW', 
  'projection.timestamp.type'='date', 
  'storage.location.template'='s3://bucket/AWSLogs/account-id/CloudTrail/aws-region/${timestamp}')

select sum(b.statistics.datascannedinbytes),
a.userIdentity.username 
from 
cloudtrail_logs a,
athena_queries b 
where 
a.eventsource='athena.amazonaws.com' and a.eventName='StartQueryExecution' and 
a.responseElements != 'null' and substr(a.responseElements,22,36) = b.queryexecutionid 
group by 
a.userIdentity.username