sql/moz-fx-data-shared-prod/cloudflare_derived/os_usage_v1/query.py

# Load libraries
import json
import os
from datetime import datetime, timedelta

import pandas as pd
import requests
from argparse import ArgumentParser
from google.cloud import bigquery
from google.cloud import storage

# Configurations
os_usg_configs = {
    "timeout_limit": 500,
    "device_types": ["DESKTOP", "MOBILE", "OTHER", "ALL"],
    "locations": [
        "ALL",
        "AU",
        "BE",
        "BG",
        "BR",
        "CA",
        "CH",
        "CZ",
        "DE",
        "DK",
        "EE",
        "ES",
        "FI",
        "FR",
        "GB",
        "HR",
        "IE",
        "IN",
        "IT",
        "JP",
        "KE",
        "CY",
        "LV",
        "LT",
        "LU",
        "HU",
        "MT",
        "MX",
        "NL",
        "AT",
        "PL",
        "PT",
        "RO",
        "SI",
        "SK",
        "US",
        "SE",
        "GR",
    ],
    "bucket": "gs://moz-fx-data-prod-external-data/",
    "bucket_no_gs": "moz-fx-data-prod-external-data",
    "results_stg_gcs_fpth": "cloudflare/os_usage/RESULTS_STAGING/%s_results_from_%s_run.csv",
    "results_archive_gcs_fpth": "cloudflare/os_usage/RESULTS_ARCHIVE/%s_results_from_%s_run.csv",
    "errors_stg_gcs_fpth": "cloudflare/os_usage/ERRORS_STAGING/%s_results_from_%s_run.csv",
    "errors_archive_gcs_fpth": "cloudflare/os_usage/ERRORS_ARCHIVE/%s_results_from_%s_run.csv",
    "gcp_project_id": "moz-fx-data-shared-prod",
    "gcp_conn_id": "google_cloud_shared_prod",
    "results_bq_stg_table": "moz-fx-data-shared-prod.cloudflare_derived.os_results_stg",
    "errors_bq_stg_table": "moz-fx-data-shared-prod.cloudflare_derived.os_errors_stg",
}

#Load the Cloudflare API Token
cloudflare_api_token = os.getenv("CLOUDFLARE_AUTH_TOKEN")

# Define a function to move a GCS object then delete the original
def move_blob(bucket_name, blob_name, destination_bucket_name, destination_blob_name):
    """Moves a blob from one bucket to another with a new name."""
    storage_client = storage.Client()
    source_bucket = storage_client.bucket(bucket_name)
    source_blob = source_bucket.blob(blob_name)
    destination_bucket = storage_client.bucket(destination_bucket_name)
    destination_generation_match_precondition = None

    blob_copy = source_bucket.copy_blob(
        source_blob,
        destination_bucket,
        destination_blob_name,
        if_generation_match=destination_generation_match_precondition,
    )
    source_bucket.delete_blob(blob_name)

    print(
        "Blob {} in bucket {} moved to blob {} in bucket {}.".format(
            source_blob.name,
            source_bucket.name,
            blob_copy.name,
            destination_bucket.name,
        )
    )


# Function to configure the API URL
def generate_os_timeseries_api_call(strt_dt, end_dt, agg_int, location, device_type):
    """Generate the API call for Operating System Usage Data."""
    if location == "ALL" and device_type == "ALL":
        os_usage_api_url = f"https://api.cloudflare.com/client/v4/radar/http/timeseries_groups/os?dateStart={strt_dt}T00:00:00.000Z&dateEnd={end_dt}T00:00:00.000Z&format=json&aggInterval={agg_int}"
    elif location != "ALL" and device_type == "ALL":
        os_usage_api_url = f"https://api.cloudflare.com/client/v4/radar/http/timeseries_groups/os?dateStart={strt_dt}T00:00:00.000Z&dateEnd={end_dt}T00:00:00.000Z&location={location}&format=json&aggInterval={agg_int}"
    elif location == "ALL" and device_type != "ALL":
        os_usage_api_url = f"https://api.cloudflare.com/client/v4/radar/http/timeseries_groups/os?dateStart={strt_dt}T00:00:00.000Z&dateEnd={end_dt}T00:00:00.000Z&deviceType={device_type}&format=json&aggInterval={agg_int}"
    else:
        os_usage_api_url = f"https://api.cloudflare.com/client/v4/radar/http/timeseries_groups/os?dateStart={strt_dt}T00:00:00.000Z&dateEnd={end_dt}T00:00:00.000Z&location={location}&deviceType={device_type}&format=json&aggInterval={agg_int}"
    return os_usage_api_url


def get_os_usage_data(date_of_interest, auth_token):
    """Pull OS usage data from the Cloudflare API and save errors & results to GCS."""
    # Calculate start date and end date
    logical_dag_dt = date_of_interest
    logical_dag_dt_as_date = datetime.strptime(logical_dag_dt, "%Y-%m-%d").date()
    start_date = logical_dag_dt_as_date - timedelta(days=4)
    end_date = start_date + timedelta(days=1)
    print("Start Date: ", start_date)
    print("End Date: ", end_date)

    # Configure request headers
    bearer_string = f"Bearer {auth_token}"
    headers = {"Authorization": bearer_string}

    # Initialize the empty results & errors dataframe
    result_df = pd.DataFrame(
        {
            "Timestamps": [],
            "OS": [],
            "Location": [],
            "DeviceType": [],
            "Share": [],
            "ConfidenceLevel": [],
            "AggrInterval": [],
            "Normalization": [],
            "LastUpdatedTS": [],
        }
    )

    # Initialize an errors dataframe
    errors_df = pd.DataFrame(
        {"StartTime": [], "EndTime": [], "Location": [], "DeviceType": []}
    )

    # Go through all combinations, submit API requests
    for device_type in os_usg_configs["device_types"]:
        for loc in os_usg_configs["locations"]:
            print("Device Type: ", device_type)
            print("Loc: ", loc)

            # Generate the URL with given parameters
            os_usage_api_url = generate_os_timeseries_api_call(
                start_date, end_date, "1d", loc, device_type
            )
            try:
                # Call the API and save the response as JSON
                response = requests.get(
                    os_usage_api_url,
                    headers=headers,
                    timeout=os_usg_configs["timeout_limit"],
                )
                response_json = json.loads(response.text)

                # If response was successful, get the result
                if response_json["success"] is True:
                    result = response_json["result"]
                    # Parse metadata
                    conf_lvl = result["meta"]["confidenceInfo"]["level"]
                    aggr_intvl = result["meta"]["aggInterval"]
                    nrmlztn = result["meta"]["normalization"]
                    lst_upd = result["meta"]["lastUpdated"]
                    data_dict = result["serie_0"]

                    for key, val in data_dict.items():
                        if key != 'timestamps':
                            new_result_df = pd.DataFrame(
                                {
                                    "Timestamps": data_dict["timestamps"],
                                    "OS": [key] * len(val),
                                    "Location": [loc] * len(val),
                                    "DeviceType": [device_type] * len(val),
                                    "Share": val,
                                    "ConfidenceLevel": [conf_lvl] * len(val),
                                    "AggrInterval": [aggr_intvl] * len(val),
                                    "Normalization": [nrmlztn] * len(val),
                                    "LastUpdatedTS": [lst_upd] * len(val),
                                }
                            )
                            result_df = pd.concat([result_df, new_result_df])

                # If response was not successful, get the errors
                else:
                    # errors = response_json["errors"]  # Maybe add to capture, right now not using this
                    new_errors_df = pd.DataFrame(
                        {
                            "StartTime": [start_date],
                            "EndTime": [end_date],
                            "Location": [loc],
                            "DeviceType": [device_type],
                        }
                    )
                    errors_df = pd.concat([errors_df, new_errors_df])
            except:
                new_errors_df = pd.DataFrame(
                    {
                        "StartTime": [start_date],
                        "EndTime": [end_date],
                        "Location": [loc],
                        "DeviceType": [device_type],
                    }
                )
                errors_df = pd.concat([errors_df, new_errors_df])

    result_fpath = os_usg_configs["bucket"] + os_usg_configs["results_stg_gcs_fpth"] % (
        start_date,
        logical_dag_dt,
    )
    errors_fpath = os_usg_configs["bucket"] + os_usg_configs["errors_stg_gcs_fpth"] % (
        start_date,
        logical_dag_dt,
    )

    result_df.to_csv(result_fpath, index=False)
    errors_df.to_csv(errors_fpath, index=False)
    print("Wrote errors to: ", errors_fpath)
    print("Wrote results to: ", result_fpath)

    # Write a summary to the logs
    len_results = str(len(result_df))
    len_errors = str(len(errors_df))
    result_summary = f"# Result Rows: {len_results}; # of Error Rows: {len_errors}"
    return result_summary


def main():
    """Call the API, save data to GCS, load to BQ staging, delete & load to BQ gold"""
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--date", required=True)
    parser.add_argument("--cloudflare_api_token", default=cloudflare_api_token)
    parser.add_argument("--project", default="moz-fx-data-shared-prod")
    parser.add_argument("--dataset", default="cloudflare_derived")

    args = parser.parse_args()
    print("Running for date: ")
    print(args.date)

    # STEP 1 - Pull the data from the API, save results & errors to GCS staging area
    result_summary = get_os_usage_data(args.date, args.cloudflare_api_token)
    print("result_summary")
    print(result_summary)

    # Create a bigquery client
    client = bigquery.Client(args.project)

    result_uri = os_usg_configs["bucket"] + os_usg_configs["results_stg_gcs_fpth"] % (
        datetime.strptime(args.date, "%Y-%m-%d").date() - timedelta(days=4),
        args.date,
    )
    error_uri = os_usg_configs["bucket"] + os_usg_configs["errors_stg_gcs_fpth"] % (
        datetime.strptime(args.date, "%Y-%m-%d").date() - timedelta(days=4),
        args.date,
    )
    print("result_uri")
    print(result_uri)

    print("error_uri")
    print(error_uri)

    # STEP 2 - Copy the result data from GCS staging to BQ staging table
    load_result_csv_to_bq_stg_job = client.load_table_from_uri(
        result_uri,
        os_usg_configs["results_bq_stg_table"],
        job_config=bigquery.LoadJobConfig(
            create_disposition="CREATE_IF_NEEDED",
            write_disposition="WRITE_TRUNCATE",
            schema=[
                {"name": "Timestamps", "type": "TIMESTAMP", "mode": "REQUIRED"},
                {"name": "OS", "type": "STRING", "mode": "NULLABLE"},
                {"name": "Location", "type": "STRING", "mode": "NULLABLE"},
                {"name": "DeviceType", "type": "STRING", "mode": "NULLABLE"},
                {"name": "Share", "type": "NUMERIC", "mode": "NULLABLE"},
                {"name": "ConfidenceLevel", "type": "STRING", "mode": "NULLABLE"},
                {"name": "AggrInterval", "type": "STRING", "mode": "NULLABLE"},
                {"name": "Normalization", "type": "STRING", "mode": "NULLABLE"},
                {"name": "LastUpdatedTS", "type": "TIMESTAMP", "mode": "NULLABLE"},
            ],
            skip_leading_rows=1,
            source_format=bigquery.SourceFormat.CSV,
        ),
    )

    load_result_csv_to_bq_stg_job.result()
    result_bq_stg_tbl = client.get_table(os_usg_configs["results_bq_stg_table"])
    print("Loaded {} rows to results staging.".format(result_bq_stg_tbl.num_rows))

    # STEP 3 - Copy the error data from GCS staging to BQ staging table
    load_error_csv_to_bq_stg_job = client.load_table_from_uri(
        error_uri,
        os_usg_configs["errors_bq_stg_table"],
        job_config=bigquery.LoadJobConfig(
            create_disposition="CREATE_IF_NEEDED",
            write_disposition="WRITE_TRUNCATE",
            schema=[
                {"name": "StartDate", "type": "DATE", "mode": "REQUIRED"},
                {"name": "EndDate", "type": "DATE", "mode": "REQUIRED"},
                {"name": "Location", "type": "STRING", "mode": "NULLABLE"},
                {"name": "DeviceType", "type": "STRING", "mode": "NULLABLE"},
            ],
            skip_leading_rows=1,
            source_format=bigquery.SourceFormat.CSV,
        ),
    )

    load_error_csv_to_bq_stg_job.result()
    error_bq_stg_tbl = client.get_table(os_usg_configs["errors_bq_stg_table"])
    print("Loaded {} rows to errors staging.".format(error_bq_stg_tbl.num_rows))

    # STEP 4 - Delete results from gold for this day, if there are any already (so if rerun, no issues will occur)
    del_exstng_gold_res_for_date = f"""DELETE FROM `moz-fx-data-shared-prod.cloudflare_derived.os_usage_v1` WHERE dte = DATE_SUB('{args.date}', INTERVAL 4 DAY)  """
    del_gold_res_job = client.query(del_exstng_gold_res_for_date)
    del_gold_res_job.result()
    print("Deleted anything already existing for this date from results gold")

    # STEP 5 - Delete errors from gold for this day, if there are any already (so if rerun, no issues will occur)
    del_exstng_gold_err_for_date = f"""DELETE FROM `moz-fx-data-shared-prod.cloudflare_derived.os_usage_errors_v1` WHERE dte = DATE_SUB('{args.date}', INTERVAL 4 DAY) """
    del_gold_err_job = client.query(del_exstng_gold_err_for_date)
    del_gold_err_job.result()
    print("Deleted anything already existing for this date from errors gold")

    # STEP 6 - Load results from stage to gold # NEED TO UPDATE THIS STILL
    os_usg_stg_to_gold_query = f""" INSERT INTO `moz-fx-data-shared-prod.cloudflare_derived.os_usage_v1`
SELECT 
CAST(Timestamps AS date) AS dte,
OS AS os,
Location AS location,
DeviceType AS device_type,
Share AS os_share,
Normalization AS normalization_type,
LastUpdatedTS AS last_updated_ts
FROM `moz-fx-data-shared-prod.cloudflare_derived.os_results_stg`
WHERE CAST(Timestamps as date) = DATE_SUB('{args.date}', INTERVAL 4 DAY) """
    load_res_to_gold = client.query(os_usg_stg_to_gold_query)
    load_res_to_gold.result()

    # STEP 7 - Load errors from stage to gold
    os_usg_errors_stg_to_gold_query = f""" INSERT INTO `moz-fx-data-shared-prod.cloudflare_derived.os_usage_errors_v1`
SELECT
StartDate AS dte,
Location AS location,
DeviceType AS device_type
FROM `moz-fx-data-shared-prod.cloudflare_derived.os_errors_stg`
WHERE StartDate = DATE_SUB('{args.date}', INTERVAL 4 DAY) """
    load_err_to_gold = client.query(os_usg_errors_stg_to_gold_query)
    load_err_to_gold.result()

    # STEP 8 - Copy the result CSV from stage to archive, then delete from stage
    # Calculate the fpaths we will use ahead of time
    result_stg_fpath = os_usg_configs["results_stg_gcs_fpth"] % (
        datetime.strptime(args.date, "%Y-%m-%d").date() - timedelta(days=4),
        args.date,
    )
    result_archive_fpath = os_usg_configs["results_archive_gcs_fpth"] % (
        datetime.strptime(args.date, "%Y-%m-%d").date() - timedelta(days=4),
        args.date,
    )
    move_blob(
        os_usg_configs["bucket_no_gs"],
        result_stg_fpath,
        os_usg_configs["bucket_no_gs"],
        result_archive_fpath,
    )

    # STEP 9 - Copy the error CSV from stage to archive, then delete from stage
    error_stg_fpath = os_usg_configs["errors_stg_gcs_fpth"] % (
        datetime.strptime(args.date, "%Y-%m-%d").date() - timedelta(days=4),
        args.date,
    )
    error_archive_fpath = os_usg_configs["errors_archive_gcs_fpth"] % (
        datetime.strptime(args.date, "%Y-%m-%d").date() - timedelta(days=4),
        args.date,
    )
    move_blob(
        os_usg_configs["bucket_no_gs"],
        error_stg_fpath,
        os_usg_configs["bucket_no_gs"],
        error_archive_fpath,
    )


if __name__ == "__main__":
    main()