Ingesting TV Response Data into Snowflake

import boto3
import os
import snowflake.connector

from boto3.session import Session

### Semantic Version ###
# Updated 2021-03-11

VERSION = "0.1.1"

### Credentials ###

TATARI_SLUG = os.environ["TATARI_SLUG"]
AWS_ACCESS_KEY_ID = os.environ["AWS_ACCESS_KEY_ID"]
AWS_SECRET_ACCESS_KEY = os.environ["AWS_SECRET_ACCESS_KEY"]

SNOWFLAKE_USER = os.environ['SNOWFLAKE_USER']
SNOWFLAKE_PASSWORD = os.environ['SNOWFLAKE_PASSWORD']
SNOWFLAKE_ACCOUNT_ID = os.environ["SNOWFLAKE_ACCOUNT_ID"]

### Configuration ##
# Update these as desired to work within your data pipeline
# Stage and table names will be {prefix}_{report}; ex: raw.public.tatari_splot_level_metrics

PREFIX = "TATARI.DATA"
DATABASE = "YOUR_DB"
WAREHOSUE = "YOUR_WH"

# These are the Tatari reports you'll be importing. Remove the ones you don't have or are uninterested in.

REPORTS = [
    "linear_spend_and_impressions",
    "spot_level_metrics",
    "streaming_spend_and_impressions",
    "user_attribution_view_through_streaming",
    "user_attribution_view_through_streaming_apps",
    "streaming_metrics"
]

### Consts ###
# These should not be changed. If you need to change these, check in with Tatari Support

ASSUME_ROLE_SESSION = "snowflake-import"
ASSUME_ROLE_ARN = f"arn:aws:iam::878256633362:role/ClientRoles/DataShare/ClientRole-DataShare-{TATARI_SLUG}"

### Schema ##
# Column types can be modified, but not column names (they are used to detect header rows)
# or column order (needs to match up with the CSV). If type is unspecified, will be VARCHAR

SCHEMA = {
    "linear_spend_and_impressions": [
        "filename VARCHAR",
        "spot_datetime",
        "creative_name",
        "creative_code",
        "network",
        "program",
        "spend",
        "impressions",
        "lift",
        "is_deadzoned",
    ],

    "spot_level_metrics": [
        "filename VARCHAR",
        "spot_id",
        "broadcast_week",
        "airdate",
        "creative_name",
        "creative_code",
        "network",
        "rotation",
        "spend",
        "conversion_metric",
        "lift",
    ],
    "streaming_spend_and_impressions": [
        "filename VARCHAR",
        "tatari_id",
        "user_id",
        "event_datetime",
        "creative_code",
        "event_type",
        "impression_datetime",
        "platform",
        "campaign_id",
    ],
    "user_attribution_view_through_streaming": [
        "filename VARCHAR",
        "tatari_id",
        "user_id",
        "event_datetime",
        "creative_code",
        "event_type",
        "impression_datetime",
        "platform",
        "campaign_id",
    ],
    "user_attribution_view_through_streaming_apps": [
        "filename VARCHAR",
        "impression_dt",
        "creative_code",
        "campaign_id",
        "event_type",
        "event_dt",
        "device_id",
        "publisher",
    ],
    "streaming_metrics": [
        "filename VARCHAR",
        "campaign_id",
        "impression_date",
        "creative_code",
        "publisher",
        "spend",
        "methodology",
        "conversion_metric",
        "lift"
    ]
}

##################################################

def schema_string_for(report):
    # Defaults cols to VARCHAR if no type is specified
    return ", ".join(
        [
            f"{col.split(' ')[0]} {(col + ' VARCHAR').split(' ')[1]}"
            for col in SCHEMA[report]
        ]
    )

def get_assumed_credentials():
    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=ASSUME_ROLE_ARN, RoleSessionName=ASSUME_ROLE_SESSION
    )

    return dict(
        aws_access_key_id=response["Credentials"]["AccessKeyId"],
        aws_secret_access_key=response["Credentials"]["SecretAccessKey"],
        aws_session_token=response["Credentials"]["SessionToken"],
    )

def stage_from_s3(report, credentials):
    return """\
 CREATE or REPLACE TEMPORARY STAGE {PREFIX}.stage_{report}
 URL = 's3://tatari-reports-exports/{SLUG}/{report}/v2/'
 CREDENTIALS = ( AWS_KEY_ID='{aws_access_key_id}' AWS_SECRET_KEY='{aws_secret_access_key}' AWS_TOKEN='{aws_session_token}' )
 """.format(
        SLUG=TATARI_SLUG, PREFIX=PREFIX, report=report, **credentials
    )

def ensure_table_for(report):
    return (
        f"CREATE or REPLACE TABLE {PREFIX}.tmp_{report} ({schema_string_for(report)})"
    )

def copy_from_stage_to_table(report):
    # Schema --> "t.$1, t.$2, ..."
    colSplat = ", ".join([f"t.${x+1}" for x in range(len(SCHEMA[report]) - 1)])

    # Grabs the name of the first column of the report
    # linear_spend_and_impressions --> spot_datetime
    colName = SCHEMA[report][1].split(" ")[0]

    # Copies all the data from the stage into the table,
    # using colName to detect and skip the header rows
    return """\
  INSERT INTO {PREFIX}.tmp_{report}
  SELECT metadata$filename AS filename, {colSplat} 
  FROM @{PREFIX}.stage_{report} t
  WHERE t.$1 != '{colName}'
  AND metadata$filename NOT IN (select distinct(filename) from TATARI.DATA.tmp_linear_spend_and_impressions)
 """.format(
        report=report, PREFIX=PREFIX, colSplat=colSplat, colName=colName
    )

##################################################

### Execution ###

logging = True

if __name__ == "__main__":
    if logging:
        print("connecting to snowflake")

    connection = snowflake.connector.connect(
        user=SNOWFLAKE_USER, password=SNOWFLAKE_PASSWORD, account=SNOWFLAKE_ACCOUNT_ID
    )

    if logging:
        print("getting AWS credentials")

    credentials = get_assumed_credentials()
    cursor = connection.cursor()

    try:
        if logging:
            print("entering report loop")
        for report in REPORTS:
            sql = [
                f"use #{DATABASE}",
                f"use WAREHOUSE #{WAREHOUSE}",
                stage_from_s3(report, credentials),
                ensure_table_for(report),
                copy_from_stage_to_table(report),
            ]

            if logging:
                print(report)
                print(";\n\n".join(sql))

            [cursor.execute(statement) for statement in sql]

    finally:
        cursor.close()
    connection.close()

Dockerfile