mozetl/taar/taar_dynamo.py

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""
This module replicates the scala script over at

https://github.com/mozilla/telemetry-batch-view/blob/1c544f65ad2852703883fe31a9fba38c39e75698/src/main/scala/com/mozilla/telemetry/views/HBaseAddonRecommenderView.scala

This should be invoked with something like this:

spark-submit \
    --master=spark://ec2-52-32-39-246.us-west-2.compute.amazonaws.com taar_dynamo.py \
    --date=20180218  \
    --region=us-west-2  \
    --table=taar_addon_data_20180206 \
    --prod-iam-role=arn:aws:iam::361527076523:role/taar-write-dynamodb-from-dev

"""

from datetime import date
from datetime import datetime
from datetime import timedelta
from pprint import pprint
from pyspark import SparkConf
from pyspark.sql import SparkSession
from pyspark.sql import Window
from pyspark.sql.functions import desc, row_number
import click
import dateutil.parser
import json
import boto3
from boto3.dynamodb.types import Binary as DynamoBinary
import time
import zlib
from .taar_utils import hash_telemetry_id

# We use the os and threading modules to generate a spark worker
# specific identity:w

import os
import threading

MAX_RECORDS = 200
EMPTY_TUPLE = (0, 0, [], [])


class CredentialSingleton:
    def __init__(self):
        self._credentials = None
        self._lock = threading.RLock()

    def __getstate__(self):
        return {"credentials": self._credentials}

    def __setstate__(self, state):
        # This is basically the constructor all over again
        self._credentials = state["credentials"]
        self._lock = threading.RLock()

    def getInstance(self, prod_iam_role):
        with self._lock:
            # If credentials exist, make sure we haven't expire them yet
            if self._credentials is not None:

                # Credentials should expire if the expiry time is sooner
                # than the next 5 minutes
                five_minute_from_now = datetime.now() + timedelta(minutes=5)
                if self._credentials["expiry"] <= five_minute_from_now:
                    self._credentials = None

            if self._credentials is None:
                self._credentials = self.get_new_creds(prod_iam_role)

            return self._credentials["cred_args"]

    def get_new_creds(self, prod_iam_role):
        client = boto3.client("sts")
        session_name = "taar_dynamo_%s_%s" % (
            os.getpid(),
            threading.current_thread().ident,
        )

        # 30 minutes to flush 50 records should be ridiculously
        # generous
        response = client.assume_role(
            RoleArn=prod_iam_role, RoleSessionName=session_name, DurationSeconds=60 * 30
        )

        raw_creds = response["Credentials"]
        cred_args = {
            "aws_access_key_id": raw_creds["AccessKeyId"],
            "aws_secret_access_key": raw_creds["SecretAccessKey"],
            "aws_session_token": raw_creds["SessionToken"],
        }

        # Set the expiry of this credential to be 30 minutes
        return {
            "expiry": datetime.now() + timedelta(minutes=30),
            "cred_args": cred_args,
        }


def json_serial(obj):
    """JSON serializer for objects not serializable by default json code"""

    try:
        if isinstance(obj, (datetime, date)):
            return obj.isoformat()
    except Exception:
        # Some dates are invalid and won't serialize to
        # ISO format if the year is < 1601.  Yes. This actually
        # happens.  Force the date to epoch in this case
        return date(1970, 1, 1).isoformat()
    raise TypeError("Type %s not serializable" % type(obj))


def filterDateAndClientID(row_jstr):
    """
    Filter out any rows where the client_id is None or where the
    subsession_start_date is not a valid date
    """
    (row, jstr) = row_jstr
    try:
        assert row.client_id is not None
        assert row.client_id != ""
        some_date = dateutil.parser.parse(row.subsession_start_date)
        if some_date.year < 1970:
            return False
        return True
    except Exception:
        return False


def list_transformer(row_jsonstr):
    """
    We need to merge two elements of the row data - namely the
    client_id and the start_date into the main JSON blob.

    This is then packaged into a 4-tuple of :

    The first integer represents the number of records that have been
    pushed into DynamoDB.

    The second is the length of the JSON data list. This prevents us
    from having to compute the length of the JSON list unnecessarily.

    The third element of the tuple is the list of JSON data.

    The fourth element is a list of invalid JSON blobs.  We maintain
    this to be no more than 50 elements long.
    """
    (row, json_str) = row_jsonstr
    client_id = row.client_id
    start_date = dateutil.parser.parse(row.subsession_start_date)
    start_date = start_date.date()
    start_date = start_date.strftime("%Y%m%d")
    jdata = json.loads(json_str)
    jdata["client_id"] = client_id
    jdata["start_date"] = start_date

    # Filter out keys with an empty value
    jdata = {key: value for key, value in list(jdata.items()) if value}

    # We need to return a 4-tuple of values
    # (numrec_dynamodb_pushed, json_list_length, json_list, error_json)

    # These 4-tuples can be reduced in a map/reduce
    return (0, 1, [jdata], [])


# TODO: singletons are hard to test - this should get refactored into
# some kind of factory or injectable dependency to the DynamoReducer
# so that we can mock out the singleton
credentials = CredentialSingleton()


class DynamoReducer(object):
    def __init__(self, prod_iam_role, region_name=None, table_name=None):

        if region_name is None:
            region_name = "us-west-2"

        if table_name is None:
            table_name = "taar_addon_data"

        self._region_name = region_name
        self._table_name = table_name
        self._prod_iam_role = prod_iam_role

    def hash_client_ids(self, data_tuple):
        """
        # Clobber the client_id by using sha256 hashes encoded as hex
        # Based on the js code in Fx

        """

        for item in data_tuple[2]:
            client_id = item["client_id"]
            item["client_id"] = hash_telemetry_id(client_id)

    def push_to_dynamo(self, data_tuple):  # noqa
        """
        This connects to DynamoDB and pushes records in `item_list` into
        a table.

        We accumulate a list of up to 50 elements long to allow debugging
        of write errors.
        """
        # Transform the data into something that DynamoDB will always
        # accept
        # Set TTL to 60 days from now
        ttl = int(time.time()) + 60 * 60 * 24 * 60

        self.hash_client_ids(data_tuple)

        item_list = [
            {
                "client_id": item["client_id"],
                "TTL": ttl,
                "json_payload": DynamoBinary(
                    zlib.compress(json.dumps(item, default=json_serial).encode("utf8"))
                ),
            }
            for item in data_tuple[2]
        ]

        # Obtain credentials from the singleton

        print("Using prod iam role: %s" % self._prod_iam_role)
        if self._prod_iam_role is not None:
            cred_args = credentials.getInstance(self._prod_iam_role)
        else:
            cred_args = {}

        conn = boto3.resource("dynamodb", region_name=self._region_name, **cred_args)

        table = conn.Table(self._table_name)
        try:
            with table.batch_writer(overwrite_by_pkeys=["client_id"]) as batch:
                for item in item_list:
                    batch.put_item(Item=item)
            return []
        except Exception:
            # Something went wrong with the batch write write.
            if len(data_tuple[3]) == 50:
                # Too many errors already accumulated, just short circuit
                # and return
                return []
            try:
                error_accum = []
                conn = boto3.resource("dynamodb", region_name=self._region_name)
                table = conn.Table(self._table_name)
                for item in item_list:
                    try:
                        table.put_item(Item=item)
                    except Exception:
                        error_accum.append(item)
                return error_accum
            except Exception:
                # Something went wrong with the entire DynamoDB
                # connection. Just return the entire list of
                # JSON items
                return item_list

    def dynamo_reducer(self, list_a, list_b, force_write=False):
        """
        This function can be used to reduce tuples of the form in
        `list_transformer`. Data is merged and when MAX_RECORDS
        number of JSON blobs are merged, the list of JSON is batch written
        into DynamoDB.
        """
        new_list = [
            list_a[0] + list_b[0],
            list_a[1] + list_b[1],
            list_a[2] + list_b[2],
            list_a[3] + list_b[3],
        ]

        if new_list[1] >= MAX_RECORDS or force_write:
            error_blobs = self.push_to_dynamo(new_list)
            if len(error_blobs) > 0:
                # Gather up to maximum 50 error blobs
                new_list[3].extend(error_blobs[: 50 - new_list[1]])
                # Zero out the number of accumulated records
                new_list[1] = 0
            else:
                # No errors during write process
                # Update number of records written to dynamo
                new_list[0] += new_list[1]
                # Zero out the number of accumulated records
                new_list[1] = 0
                # Clear out the accumulated JSON records
                new_list[2] = []

        return tuple(new_list)


def etl(spark, run_date, region_name, table_name, prod_iam_role, sample_rate):
    """
    This function is responsible for extract, transform and load.

    Data is extracted from Parquet files in Amazon S3.
    Transforms and filters are applied to the data to create
    3-tuples that are easily merged in a map-reduce fashion.

    The 3-tuples are then loaded into DynamoDB using a map-reduce
    operation in Spark.
    """

    rdd = extract_transform(spark, run_date, sample_rate)
    result = load_rdd(prod_iam_role, region_name, table_name, rdd)
    return result


def extract_transform(spark, run_date, sample_rate=0):
    currentDate = run_date
    currentDateString = currentDate.strftime("%Y%m%d")
    print("Processing %s" % currentDateString)

    # Get the data for the desired date out of parquet
    template = "s3://telemetry-parquet/main_summary/v4/submission_date_s3=%s"
    datasetForDate = spark.read.parquet(template % currentDateString)

    if sample_rate is not None and sample_rate != 0:
        print("Sample rate set to %0.9f" % sample_rate)
        datasetForDate = datasetForDate.sample(False, sample_rate)
    else:
        print("No sampling on dataset")

    print("Parquet data loaded")

    # Get the most recent (client_id, subsession_start_date) tuple
    # for each client since the main_summary might contain
    # multiple rows per client. We will use it to filter out the
    # full table with all the columns we require.

    clientShortList = datasetForDate.select(
        "client_id",
        "subsession_start_date",
        row_number()
        .over(Window.partitionBy("client_id").orderBy(desc("subsession_start_date")))
        .alias("clientid_rank"),
    )
    print("clientShortList selected")
    clientShortList = clientShortList.where("clientid_rank == 1").drop("clientid_rank")
    print("clientShortList selected")

    select_fields = [
        "client_id",
        "subsession_start_date",
        "subsession_length",
        "city",
        "locale",
        "os",
        "places_bookmarks_count",
        "scalar_parent_browser_engagement_tab_open_event_count",
        "scalar_parent_browser_engagement_total_uri_count",
        "scalar_parent_browser_engagement_unique_domains_count",
        "active_addons",
        "disabled_addons_ids",
    ]
    dataSubset = datasetForDate.select(*select_fields)
    print("datasetForDate select fields completed")

    # Join the two tables: only the elements in both dataframes
    # will make it through.
    clientsData = dataSubset.join(
        clientShortList, ["client_id", "subsession_start_date"]
    )

    print("clientsData join with client_id and subsession_start_date")

    # Convert the DataFrame to JSON and get an RDD out of it.
    subset = clientsData.select("client_id", "subsession_start_date")

    print("clientsData select of client_id and subsession_start_date completed")

    jsonDataRDD = clientsData.select(
        "city",
        "subsession_start_date",
        "subsession_length",
        "locale",
        "os",
        "places_bookmarks_count",
        "scalar_parent_browser_engagement_tab_open_event_count",
        "scalar_parent_browser_engagement_total_uri_count",
        "scalar_parent_browser_engagement_unique_domains_count",
        "active_addons",
        "disabled_addons_ids",
    ).toJSON()

    print("jsonDataRDD selected")
    rdd = subset.rdd.zip(jsonDataRDD)
    print("subset rdd has been zipped")

    # Filter out any records with invalid dates or client_id
    filtered_rdd = rdd.filter(filterDateAndClientID)
    print("rdd filtered by date and client_id")

    # Transform the JSON elements into a 4-tuple as per docstring
    merged_filtered_rdd = filtered_rdd.map(list_transformer)
    print("rdd has been transformed into tuples")

    return merged_filtered_rdd


def load_rdd(prod_iam_role, region_name, table_name, rdd):
    # Apply a MapReduce operation to the RDD
    dynReducer = DynamoReducer(prod_iam_role, region_name, table_name)

    reduction_output = rdd.reduce(dynReducer.dynamo_reducer)
    print("1st pass dynamo reduction completed")

    # Apply the reducer one more time to force any lingering
    # data to get pushed into DynamoDB.
    final_reduction_output = dynReducer.dynamo_reducer(
        reduction_output, EMPTY_TUPLE, force_write=True
    )
    return final_reduction_output


def run_etljob(spark, run_date, region_name, table_name, prod_iam_role, sample_rate):
    reduction_output = etl(
        spark, run_date, region_name, table_name, prod_iam_role, sample_rate
    )
    report_data = (reduction_output[0], reduction_output[1])
    print("=" * 40)
    print(
        "%d records inserted to DynamoDB.\n%d records remaining in queue." % report_data
    )
    print("=" * 40)
    return reduction_output


@click.command()
@click.option("--date", required=True)  # YYYYMMDD
@click.option("--region", default="us-west-2")
@click.option("--table", default="taar_addon_data_20180206")
@click.option(
    "--prod-iam-role",
    default="arn:aws:iam::361527076523:role/taar-write-dynamodb-from-dev",
)
@click.option("--sample-rate", default=0)
def main(date, region, table, prod_iam_role, sample_rate):
    APP_NAME = "HBaseAddonRecommenderView"
    conf = SparkConf().setAppName(APP_NAME)
    spark = SparkSession.builder.config(conf=conf).getOrCreate()
    date_obj = datetime.strptime(date, "%Y%m%d")

    if prod_iam_role.strip() == "":
        prod_iam_role = None

    reduction_output = run_etljob(
        spark, date_obj, region, table, prod_iam_role, sample_rate
    )
    pprint(reduction_output)