<commit_before>#!/usr/bin/env python3
# -*- encoding: utf-8
"""
Create a request to archive a bag

Usage: trigger_archive_bag.py <BAG>... [--bucket=<BUCKET_NAME>] [--topic=<TOPIC_NAME>|--api=<API>] [--sns=(true|false)]
       trigger_archive_bag.py -h | --help

Options:
  sns                    Send directly to SNS rather than through the API
                         [default: false]
  --bucket=<BUCKET_NAME> The S3 bucket containing the bags.
                         [default: wellcomecollection-assets-archive-ingest]
  --topic=<TOPIC_NAME>   The archivist topic.
                         [default: archive-storage_archivist]
   --api=<API>           The API endpoint to use
                         [default: http://api.wellcomecollection.org/prod/storage/v1/ingests]
   --sns=(true|false)    Send directly to SNS rather than through the API
                         [default: false]
  -h --help              Print this help message
"""

import boto3
import docopt
import json
import requests


def archive_bag_sns_messages(bags, bucket):
    """
    Generates bag archive messages.
    """
    for bag in bags:
        yield {
            'namespace': bucket,
            'key': bag
        }


def archive_bag_api_messages(bags, bucket):
    """
    Generates bag archive messages.
    """
    for bag in bags:
        yield {
            'uploadUrl': f"s3://{bucket}/{bag}"
        }


def build_topic_arn(topic_name):
    """Given a topic name, return the topic ARN."""
    # https://stackoverflow.com/a/37723278/1558022
    sts_client = boto3.client('sts')
    account_id = sts_client.get_caller_identity().get('Account')

    return f'arn:aws:sns:eu-west-1:{account_id}:{topic_name}'


def publish_messages(topic_arn, messages):
    """Publish a sequence of messages to an SNS topic."""
    sns_client = boto3.client('sns')
    for m in messages:
        message_as_json = json.dumps(m)
        response = sns_client.publish(
            TopicArn=topic_arn,
            MessageStructure='json',
            Message=json.dumps({
                'default': message_as_json
            }),
            Subject=f'Source: {__file__}'
        )
        response_status = response['ResponseMetadata']['HTTPStatusCode']
        print(f'{message_as_json} -> {topic_arn} [{response_status}]')
        assert response_status == 200, response


def publish_to_sns(bucket_name, bags, topic_name):
    topic_arn = build_topic_arn(topic_name)

    publish_messages(
        topic_arn=topic_arn,
        messages=archive_bag_sns_messages(bags, bucket_name)
    )


def call_ingest_api(bucket_name, bags, api):
    session = requests.Session()
    for message in archive_bag_api_messages(bags, bucket_name):
        response = session.post(api, json=message)
        print(f'{message} -> {api} [{response.status_code}] {response.json()}')


def main():
    args = docopt.docopt(__doc__)
    bags = args['<BAG>']
    bucket_name = args['--bucket']
    use_sns_directly = args['--sns']<commit_msg>update trigger_archive_bag to match API ingest request model<commit_after>#!/usr/bin/env python3
# -*- encoding: utf-8
"""
Create a request to archive a bag

Usage: trigger_archive_bag.py <BAG>... [--bucket=<BUCKET_NAME>] [--topic=<TOPIC_NAME>|--api=<API>] [--sns=(true|false)]
       trigger_archive_bag.py -h | --help

Arguments:
    BAG                    BagIt files to ingest

Examples:
    trigger_archive_bag.py b22454408.zip

Options:
    sns                    Send directly to SNS rather than through the API
                           [default: false]
    --bucket=<BUCKET_NAME> The S3 bucket containing the bags.
                           [default: wellcomecollection-assets-archive-ingest]
    --topic=<TOPIC_NAME>   The archivist topic.
                           [default: archive-storage_archivist]
    --api=<API>            The API endpoint to use
                           [default: http://api.wellcomecollection.org/storage/v1/ingests]
    --sns=(true|false)     Send directly to SNS rather than through the API
                           [default: false]
    -h --help              Print this help message
"""

import boto3
import docopt
import json
import requests


def archive_bag_sns_messages(bags, bucket):
    """
    Generates bag archive messages.
    """
    for bag in bags:
        yield {
            'namespace': bucket,
            'key': bag
        }


def archive_bag_api_messages(bags, bucket):
    """
    Generates bag archive messages.
    """
    for bag in bags:
        yield {
            'type': 'Ingest',
            'ingestType' : {
                'id': 'create',
                'type': 'IngestType'
            },
            'uploadUrl': f's3://{bucket}/{bag}'
        }


def build_topic_arn(topic_name):
    """Given a topic name, return the topic ARN."""
    # https://stackoverflow.com/a/37723278/1558022
    sts_client = boto3.client('sts')
    account_id = sts_client.get_caller_identity().get('Account')

    return f'arn:aws:sns:eu-west-1:{account_id}:{topic_name}'


def publish_messages(topic_arn, messages):
    """Publish a sequence of messages to an SNS topic."""
    sns_client = boto3.client('sns')
    for m in messages:
        message_as_json = json.dumps(m)
        response = sns_client.publish(
            TopicArn=topic_arn,
            MessageStructure='json',
            Message=json.dumps({
                'default': message_as_json
            }),
            Subject=f'Source: {__file__}'
        )
        response_status = response['ResponseMetadata']['HTTPStatusCode']
        print(f'{message_as_json} -> {topic_arn} [{response_status}]')
        assert response_status == 200, response


def publish_to_sns(bucket_name, bags, topic_name):
    topic_arn = build_topic_arn(topic_name)

    publish_messages(
        topic_arn=topic_arn,
        messages=archive_bag_sns_messages(bags, bucket_name)
    )


def call_ingest_api(bucket_name, bags, api):
    session = requests.Session()
    for message in archive_bag_api_messages(bags, bucket_name):
        response = session.post(api, json=message)
        status_code = response.status_code
        if status_code != 202:
            print_result(f'ERROR calling {api}', response)
        else:
            print(f'{message} -> {api} [{status_code}]')
            location = response.headers.get('Location')
            ingest = session.get(location)
            if location:
                print_result(location, ingest)


def print_result(description, result):
    print(description)
    dumped_json = json.dumps(result.json(), indent=2)
    print(dumped_json)


def main():
    args = docopt.docopt(__doc__)
    bags = args['<BAG>']
    bucket_name = args['--bucket']
    use_sns_directly = args['--sns']
{'fraction_non_alphanumeric': 0.11858118135945385, 'fraction_numerical': 0.009795191451469279, 'mean_word_length': 2.381334671349724, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 23, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5196723', 'n_tokens_mistral': 2111, 'n_tokens_neox': 2030, 'n_words': 457}