sync_ensemble.py

import itertools
import io
import os

import boto3
import requests

S3_BUCKET = os.getenv("S3_BUCKET", "net-mozaws-dev-us-east-1-data-ensemble")
S3_PREFIX = "datasets"

DASHBOARDS = {
    "desktop": {
        "user-activity": "https://ensemble-transposer.herokuapp.com/datasets/user-activity",
        "usage-behavior": "https://ensemble-transposer.herokuapp.com/datasets/usage-behavior",
        "hardware": "https://ensemble-transposer.herokuapp.com/datasets/hardware",
    },
}


class RequestError(Exception):
    pass


def read_dashboard(url):
    data = requests.get(url).json()
    dashboard = {}
    for product in itertools.product(data["categories"], data["metrics"]):
        key = "/".join(product)
        val = "{}/{}".format(url, key)
        dashboard[key] = val
    return dashboard


def get_and_push(bucket, key, url):
    r = requests.get(url)
    if r.status_code != 200:
        raise RequestError("response is not 200")
    client = boto3.client("s3")
    client.upload_fileobj(io.BytesIO(r.content), bucket, key,
                          ExtraArgs={"ContentType": "application/json"})


def lambda_handler(event, context):
    # iterate over sections
    for section_key, section in DASHBOARDS.items():
        # iterate over dashboards
        for dashboard_key, dashboard_url in section.items():
            # sync dashboard
            dashboard_full_key = "{}/{}/{}".format(S3_PREFIX, section_key, dashboard_key)
            print("Syncing {} to s3://{}/{}".format(dashboard_url, S3_BUCKET, dashboard_full_key))
            get_and_push(S3_BUCKET, dashboard_full_key, dashboard_url)

            # read dashboard, iterate over metrics
            dashboard = read_dashboard(dashboard_url)
            for metric_key, metric_url in dashboard.items():
                # sync metric
                metric_full_key = "{}/{}".format(dashboard_full_key, metric_key)
                print("Syncing {} to s3://{}/{}".format(metric_url, S3_BUCKET, metric_full_key))
                get_and_push(S3_BUCKET, metric_full_key, metric_url)