Sample Ingestion Script for Streaming User Attribution Report (for apps)

Below is a sample python script that demonstrates how to do an initial load as well as incremental load of the streaming user attribution report (for apps). This script will assume our data sharing role using your AWS credentials, pull all necessary files from Tatari's S3 and load it into an sqlite3 table.

Before running this script, you will need to replace the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY variables with the credentials from your AWS account that has permissions to assume this role. Note that if you are trying to use your own AWS role to assume our role, you will need to alter the code to assume role twice.

You can run this script with the following command:

python3 {script_file_path} {slug} -s {start_date} -e {end_date}

Start and end dates are optional, and will default to 60 days ago and today respectively if omitted. The script will then grab the necessary s3 keys for a report, parse them and load them into a sqlite3 table.

#!/usr/bin/env python

import boto3
import csv
import gzip
import logging
import pytz
import sqlite3

from argh import arg, dispatch_command
from boto3.session import Session
from contextlib import contextmanager
from datetime import date, datetime, timedelta
from typing import Optional

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(asctime)s: %(message)s")
BUCKET = "tatari-reports-exports"

# Replace the following with credentials of the user/role setup to access Tatari's S3
AWS_ACCESS_KEY_ID = 'Access Key ID'
AWS_SECRET_ACCESS_KEY = 'Secret Access Key ID'

################## Per-Report Configuration ##################
REPORT = "user_attribution_streaming_app"
KEY_TEMPLATE = '{0}/user_attribution_view_through_streaming_app/v2/{0}_{1}_{2}-{3}.csv.gz'
COLUMNS = [
    "impression_datetime",
    "creative_code",
    "campaign_id",
    "event_type",
    "event_datetime",
    "device_id",
    "publisher"
]

# This is the key we'll use to delete as we ingest more-mature reports
PRIMARY_COL = "impression_datetime"
PRIMARY_COL_INDEX = 0

# This is the event name we want to get total lift counts for
EVENT = 'install'

# This is used to determine offset from UTC for dates
ET = pytz.timezone('US/Eastern')


# Helper function for converting UTC date to ET datetime
def get_offset(date):
    return int(ET.localize(datetime.fromordinal(date.toordinal())).strftime('%z')[1:3])


# Different reports have different maturation cadences. We want to grab the mature reports, and then
# if the target date-range include immature data, then we also want one unreconciled report.
def derive_date_ranges(slug: str, start_date: date, end_date: date):
    """
    Returns list of keys for existing S3 reports
    that contain data between start and end date.
    """
    date_ranges = []
    one_day = timedelta(1)
    reconciled_period = timedelta(29)
    end_of_reconciled_period = date.today() - reconciled_period

    current_date = start_date
    end = min(end_date, end_of_reconciled_period)

    # Get reconciled (mature) report key(s) from start_date up to end_date
    # or end of reconciled period, whichever is earlier
    while current_date < end:
        date_ranges.append([current_date, current_date + one_day])
        current_date += one_day

    # Get one unreconciled (immature) report key if end_date is in unreconciled period
    if end_of_reconciled_period < end_date:
        date_ranges.append([current_date, current_date + reconciled_period])

    return date_ranges


def update_attribution_table(con, start_date: date, end_date: date, to_db):
    """
    Delete existing data in date range of new data and then inserts new data.
    Data in report is inclusive of start date and exclusive of end date.
    """
    # Datetimes for this report are UTC, but the data start/end is in US/Eastern
    start_offset = get_offset(start_date)
    end_offset = get_offset(end_date)

    cur = con.cursor()

    # impression_datetime range this report is +5 hours of the start/end dates of the file
    cur.execute(f'''
            DELETE FROM {REPORT}
            WHERE DATETIME({PRIMARY_COL}) >= DATETIME('{start_date} {start_offset}:00:00+00:00') 
            AND DATETIME({PRIMARY_COL}) < DATETIME('{end_date} {end_offset}:00:00+00:00');
            ''')
    con.commit()

    schema = ", ".join(COLUMNS)
    schema_values = ", ".join(["?"] * len(COLUMNS))
    cur.executemany(f"INSERT INTO {REPORT} ({schema}) VALUES ({schema_values});", to_db)
    con.commit()
    cur.close()

    return


############### Library ###############
# This code does not differ report-to-report
def get_tatari_session(slug):
    client = boto3.client(
        'sts',
        aws_access_key_id=AWS_ACCESS_KEY_ID,
        aws_secret_access_key=AWS_SECRET_ACCESS_KEY
    )
    response = client.assume_role(
        RoleArn=f'arn:aws:iam::878256633362:role/ClientRoles/DataShare/ClientRole-DataShare-{slug}',
        RoleSessionName='my-session',
    )
    return Session(
        aws_access_key_id=response['Credentials']['AccessKeyId'],
        aws_secret_access_key=response['Credentials']['SecretAccessKey'],
        aws_session_token=response['Credentials']['SessionToken'],
    )


def load_data(
    s3_client,
    con,
    slug: str,
    start: str,
    end: str
):
    """
    Recreates intended usage/ingestion of s3 reports
    """
    start_date = date.fromisoformat(start) if start else date.today() - timedelta(60)
    end_date = date.fromisoformat(end) if end else date.today()
    end_date = min(date.today(), end_date)
    date_ranges = derive_date_ranges(slug, start_date, end_date)
    logging.info(f"{len(date_ranges)} files to ingest")

    create_table(con)

    for dates in date_ranges:
        key = KEY_TEMPLATE.format(slug, REPORT, dates[0], dates[1])
        logging.info(f"Loading {key} to table")
        to_db = load_file(s3_client, key)
        if to_db:
            update_attribution_table(con, dates[0], dates[1], to_db)
    logging.info("Ingestion complete")

    return


def create_table(con):
    """
    Creates table to load report data into, truncate existing data.
    """
    schema = ", ".join(COLUMNS)
    cur = con.cursor()
    cur.execute(f"CREATE TABLE IF NOT EXISTS {REPORT} ({schema});")
    cur.execute(f"DELETE FROM {REPORT}")
    con.commit()
    cur.close()

    return


def load_file(s3_client, key):
    """
    Downloads and reads report data from S3 and returns rows to insert.
    """
    file_path = "/tmp/file.csv.gz"
    try:
        s3_client.download_file(BUCKET, key, file_path)
        to_db = read_file(file_path)
        return to_db
    except Exception as e:
        logging.error(f"Error when trying to download {key}, skipping file. Error {e}")
        return None


def read_file(file):
    """
    Returns rows to insert.
    """
    with gzip.open(file, 'rt') as f:
        reader = csv.DictReader(f)
        to_db = [tuple(r[c] for c in COLUMNS) for r in reader]

    return to_db


@contextmanager
def prepared_database(slug, start, end):
    s3_client = get_tatari_session(slug).client('s3')
    con = sqlite3.connect(f"{slug}.db")

    load_data(s3_client, con, slug, start, end)

    try:
        yield con.cursor()
    finally:
        # replace the following line with pass if you want to query the db after
        con.close()


################ Query Code ################
# This is the code we'll use to perform some simple analysis of the data
def query_sql(start, end):
    # The dashboard selector is [inclusive, inclusive],
    # but SQL checks [inclusive, exclusive)
    # so we need to tweak our date selection a little bit
    start_date = date.fromisoformat(start) if start else date.today() - timedelta(60)
    end_date = date.fromisoformat(end) + timedelta(1) if end else date.today()

    # Get time offsets between start/end dates and US/Eastern
    start_offset = get_offset(start_date)
    end_offset = get_offset(end_date)

    # impression_datetime range this report is +5 hours of the start/end dates of the file
    return(f'''
        SELECT COUNT(*) AS {EVENT}_count FROM {REPORT}
        WHERE event_type = '{EVENT}'
        AND {PRIMARY_COL} >= DATETIME('{start_date} {start_offset}:00:00+00:00') 
        AND {PRIMARY_COL} >= DATETIME('{end_date.strftime('%Y-%m-%d')} {end_offset}:00:00+00:00')
        ;''')


@arg('slug', help='This is your company\'s name within Tatari data. It\'ll appear in the URL of your dashboard')
@arg(
    '-s',
    '--start',
    help='This is the first date you\'ve selected on your dashboard. Format date as 2020-11-15.'
    'If omitted, will default to 60 days ago.'
)
@arg(
    '-e',
    '--end',
    help='This is the last date you\'ve selected on your dashboard. Format date as 2020-11-15.'
    'If omitted, will default today.'
)
def main(
    slug: str,
    start: Optional[str] = None,
    end: Optional[str] = None
):
    # Library code will now prepare the database...
    with prepared_database(slug, start, end) as db_cursor:
        # ...and now we'll run our query and get the result.
        db_cursor.execute(query_sql(start, end))
        results = db_cursor.fetchall()
        print(results)


if __name__ == '__main__':
    dispatch_command(main)