AWSTemplateFormatVersion: 2010-09-09

Description: This template deploys a SQS queue, a Lambda Function and Glue Table to consolidate S3 inventory reports.

Resources:
#region Queue
  rSqsQueue:
    Type: AWS::SQS::Queue
    Properties:
      QueueName: 'inventory-consolidation-queue'
      MessageRetentionPeriod: 1209600
      VisibilityTimeout: 600
      RedrivePolicy: { "deadLetterTargetArn" : !GetAtt rSqsDLQ.Arn , "maxReceiveCount" : 3 }


  rSqsDLQ:
    Type: AWS::SQS::Queue
    Properties:
      QueueName: 'inventory-consolidation-dlq'
      MessageRetentionPeriod: 1209600
      VisibilityTimeout: 600


  MySqsQueuePolicy:
    Type: AWS::SQS::QueuePolicy
    Properties:
      Queues:
        - !Ref rSqsQueue
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          Effect: Allow
          Principal: "*"
          Action: sqs:SendMessage
          Resource:
            Fn::GetAtt:
            - rSqsQueue
            - Arn
          Condition:
            ArnLike:
              aws:SourceArn: !Join
                - ":"
                - - "arn:aws:sns:*"
                  - !Ref "AWS::AccountId"
                  - "InventoryCentralizationTopic-*"
#endregion
#region bucket
  rInventoryReportsBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'AES256'
      BucketName: !Join
        - "-"
        - - "inventory-consolidated"
          - !Sub ${AWS::Region}
          - !Select
            - 0
            - !Split
              - "-"
              - !Select
                - 2
                - !Split
                  - "/"
                  - !Ref "AWS::StackId"
      VersioningConfiguration:
        Status: Enabled
      PublicAccessBlockConfiguration:
          BlockPublicAcls: True
          BlockPublicPolicy: True
          IgnorePublicAcls: True
          RestrictPublicBuckets: True
          
#endregion
#region processing
  rLambdaFunctionEventSourceMapping:
    Type: AWS::Lambda::EventSourceMapping
    Properties:
      BatchSize: 10
      Enabled: true
      EventSourceArn: !GetAtt  rSqsQueue.Arn
      FunctionName: !Ref rInventoryProcessingFunction

  rLambdaFunctionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: InventoryProcessing
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
        - arn:aws:iam::aws:policy/AWSXrayWriteOnlyAccess
      Path: /
      Policies:
        - PolicyName: InventoryPolicy
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action: 
                  - sqs:ReceiveMessage
                  - sqs:DeleteMessage
                  - sqs:GetQueueAttributes
                Resource:
                  - !GetAtt rSqsQueue.Arn
                  - !GetAtt rSqsDLQ.Arn
              - Effect: Allow
                Action:
                  - s3:GetObject
                Resource: 'arn:aws:s3:::inventory-*'
              - Effect: Allow
                Action:
                  - 's3:GetBucketLocation'
                Resource: '*'
              - Effect: Allow
                Action:
                  - 's3:AbortMultipartUpload'
                  - 's3:GetBucketLocation'
                  - 's3:GetObject'
                  - 's3:ListBucket'
                  - 's3:ListBucketMultipartUploads'
                  - 's3:PutObject'
                  - 's3:PutObjectAcl'
                Resource:
                  - !GetAtt rInventoryReportsBucket.Arn
                  - !Join
                    - ''
                    - - 'arn:aws:s3:::'
                      - !Ref rInventoryReportsBucket
                      - '/*'
              - Effect: Allow
                Action:
                  - 'glue:CreatePartition'
                Resource: 
                  - !Sub 'arn:aws:glue:${AWS::Region}:${AWS::AccountId}:table/${rGlueDatabase}/${rGlueTable}'
                  - !Sub 'arn:aws:glue:${AWS::Region}:${AWS::AccountId}:database/${rGlueDatabase}'
                  - !Sub 'arn:aws:glue:${AWS::Region}:${AWS::AccountId}:catalog'


  rInventoryProcessingFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      FunctionName: InventoryProcessing
      Handler: index.handler
      Code: 
        ZipFile: !Sub |
            import json
            import os
            import logging
            import boto3
            from datetime import datetime
            from boto3.s3.transfer import TransferConfig
            from botocore.exceptions import ClientError

            log = logging.getLogger(__name__)
            log.setLevel(os.getenv("LOG_LEVEL", "INFO"))

            for name in ['boto3', 's3transfer', 'botocore', 'boto']:
                logging.getLogger(name).setLevel(logging.ERROR)
            
            TARGET = "${rInventoryReportsBucket}"
            session = boto3.Session()

            def handler(event, context=None):
                log.debug(event)
                for message in event['Records']:
                    records = json.loads(json.loads(message['body'])['Message'])
                    for record in records.get('Records',[]):
                        bucket_name = record["s3"]["bucket"]["name"]
                        key = record["s3"]["object"]["key"]
                        process_inventory_file(bucket_name, key)


            def process_inventory_file(bucket_name, key, region=None):
                log.debug("Processing inventory file")
                key = key.replace('checksum', 'json')
                region = session.client('s3').get_bucket_location(Bucket=bucket_name)['LocationConstraint'] or 'us-east-1'
                resp = json.loads(get_file_content(bucket_name, key, region))
                log.debug(f"processing details : {bucket_name}, {key}, {region}")

                bucket = key.split('/')[1]
                account = key.split('/')[0]
                inventory_date = datetime.strptime(key.split('/')[-2], '%Y-%m-%dT%H-%MZ')
                date = inventory_date.strftime('%Y-%m-%d')

                for file in resp['files']:
                    src_file = file['key']
                    filename = src_file.split('/')[-1]
                    tgt_key = f"dt={date}/account={account}/region={region}/bucket_part={bucket}/{filename}"
                    log.debug("copying {} to {}".format(file['key'], TARGET))
                    copy_file(session, bucket_name, src_file, TARGET, tgt_key)
                
                add_partition(date=date, bucket_part=bucket, account=account, region=region)

            def copy_file(session, src_bucket, src_key, dst_bucket, dst_key):
                tgt = session.resource('s3').Bucket(dst_bucket)
                tgt_object = tgt.Object(dst_key)
                try:
                    tgt_object.copy(CopySource={'Bucket': src_bucket, 'Key': src_key},
                                    ExtraArgs={'ACL': 'bucket-owner-full-control'},
                                    Config=TransferConfig(5 * (1024 ** 3)), )
                except ClientError as e:
                    log.error(e)
                    raise e


            def get_file_content(bucket_name, key, region):
                s3 = session.resource("s3", region_name=region)
                o = s3.Object(bucket_name, key)
                resp = o.get()['Body'].read().decode('utf-8')
                return resp

            def add_partition(date, bucket_part, account, region):
                glue = session.client('glue')
                partition_input={}
                partition_input['StorageDescriptor'] = {
                    'Location' : f's3://{TARGET}/dt={date}/account={account}/region={region}/bucket_part={bucket_part}',
                    'InputFormat':'org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat',
                    'OutputFormat': 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat',
                    'SerdeInfo': {'SerializationLibrary': 'org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe'}
                    }
                partition_input['Values'] = [date, account, region, bucket_part]
                try :
                    glue.create_partition(
                        DatabaseName='${rGlueDatabase}',
                        TableName='${rGlueTable}',
                        PartitionInput=partition_input)
                except ClientError as e:
                    if e.response['Error']['Code'] != 'AlreadyExistsException':
                      raise e

      Runtime: python3.7
      Description: Lambda processing function to consolidate inventory reports
      MemorySize: 128
      Timeout: 600
      Role:
        Fn::GetAtt:
          - rLambdaFunctionRole
          - Arn
      Environment:
        Variables:
          PYTHONWARNINGS: ignore

  rGlueDatabase:
    Type: AWS::Glue::Database
    Properties:
      CatalogId: !Ref AWS::AccountId
      DatabaseInput:
        Description: S3 Inventory database
        Name: s3
        LocationUri: !Join ['', ['s3://', !Ref rInventoryReportsBucket]]

  rGlueTable:
    Type: AWS::Glue::Table
    Properties:
      CatalogId: !Ref AWS::AccountId
      DatabaseName: !Ref rGlueDatabase
      TableInput:
        Name: inventory
        Owner: owner
        Retention: 0
        StorageDescriptor:
          Columns:
            - Name: bucket
              Type: string
            - Name: key
              Type: string
            - Name: size
              Type: bigint
            - Name: last_modified_date
              Type: timestamp
            - Name: e_tag
              Type: string
            - Name: storage_class
              Type: string
            - Name: is_mulitpart_upload
              Type: boolean
            - Name: replication_status
              Type: string
            - Name: encryption_status
              Type: string
            - Name: object_lock_retain_until_date
              Type: timestamp
            - Name: object_lock_mode
              Type: string
            - Name: object_lock_legal_hold_status
              Type: string
            - Name: intelligent_tiering_access_tier
              Type: string
          InputFormat: org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat
          OutputFormat: org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat
          Compressed: false
          NumberOfBuckets: -1
          SerdeInfo:
            SerializationLibrary: org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe
            Parameters:
              serialization.format: '1'
          BucketColumns: []
          SortColumns: []
          StoredAsSubDirectories: false
          Location: !Sub 's3://${rInventoryReportsBucket}'
        PartitionKeys:
          - Name: dt
            Type: string
          - Name: account
            Type: string
          - Name: region
            Type: string
          - Name: bucket_part
            Type: string
        TableType: EXTERNAL_TABLE

#endregion