Google Cloud Storage

from datetime import date
import gzip
import json
import random
import time

from google.cloud import storage
import requests

# Fill this out.
PROJECT_TOKEN = ""

# Flush a batch once these limits are hit
EVENTS_PER_BATCH = 2000
BYTES_PER_BATCH = 2 * 1024 * 1024


# Adjust event timestamps of test data so that it appears recent in Mixpanel.
# NOTE: THIS IS JUST FOR DEMO PURPOSES.
# REPLACE THIS WITH YOUR OWN TRANSFORMS IN PRODUCTION.
def transform_to_event(line):
    """Convert a line of the file to a json string in Mixpanel's format."""
    event = json.loads(line)
    event["properties"]["time"] = int(date.today().strftime("%s"))
    return json.dumps(event)


def flush(batch):
    payload = gzip.compress("\n".join(batch).encode("utf-8"))
    tries = 0
    while True:
        resp = requests.post(
            "https://api.mixpanel.com/import",
            params={"strict": "1"},
            headers={
                "Content-Type": "application/x-ndjson",
                "Content-Encoding": "gzip",
                "User-Agent": "mixpanel-gcs"
            },
            auth=(PROJECT_TOKEN, ""),
            data=payload,
        )
        if resp.status_code == 429 or resp.status_code >= 500:
            time.sleep(min(2 ** tries, 60) + random.randint(1, 5))
            tries += 1
            continue
        return resp


def main(request, context):
    gcs = storage.Client()
    bucket = gcs.get_bucket(request["bucket"])
    blob = bucket.get_blob(request["name"])
    
    compressed = request["name"].endswith(".gz")

    error = None
    batch, batch_bytes = [], 0
    
    # As a convenience, check for files ending with .gz
    f = blob.open("rb")
    if request["name"].endswith(".gz") and request.get("contentEncoding") != "gzip":
        f = gzip.open(f)
    
    for line in f:
        # Here we assume each line is valid JSON representing a single Mixpanel event
        # If lines are not in Mixpanel's format, apply any transformations here.
        line = transform_to_event(line)

        batch.append(line)
        batch_bytes += len(line)

        if len(batch) == EVENTS_PER_BATCH or batch_bytes >= BYTES_PER_BATCH:
            resp = flush(batch)
            batch, batch_bytes = [], 0

            if resp.status_code != 200:
                error = resp.json()
                break

    # Flush final batch
    if batch and not error:
        resp = flush(batch)
        if resp.status_code != 200:
            error = resp.json()

    f.close()
    print(
        json.dumps(
            {
                "message": "Import complete",
                "request": request,
                "success": error is None,
                "first_error": error,
                "severity": "INFO",
            }
        )
    )

google-cloud-storage
requests