pubsub_stream.py

import csv
import json
from concurrent import futures
from google.cloud import pubsub_v1
from typing import Callable
from time import sleep
from random import random


project_id = "ab-academy-demo"
topic_id = "city_speeds"

client = pubsub_v1.PublisherClient()
topic = client.topic_path(project_id, topic_id)
publish_futures = []
filename = "city_daily_speeds-10k.csv"

counter = 0

def get_callback(
    publish_future: pubsub_v1.publisher.futures.Future, data: str
) -> Callable[[pubsub_v1.publisher.futures.Future], None]:
    def callback(publish_future: pubsub_v1.publisher.futures.Future) -> None:
        try:
            # Wait 60 seconds for the publish call to succeed.
            print(publish_future.result(timeout=60))
        except futures.TimeoutError:
            print(f"Publishing {data} timed out.")

    return callback

with open(filename, encoding="ascii", errors="surrogateescape") as data_file:
    reader = csv.reader(data_file)
    # skip the first row of the CSV file.
    next(reader)
    for line in reader:
        # prepare each line as a JSON record to match the Avro schema
        # we don't actually need the schema in our script, unless we want
        # to use a DatumWriter for binary insertion
        record = {
            "country": line[0],
            "country_code": line[1],
            "region": line[2],
            "region_code": line[3],
            "city": line[4],
            "date": line[5],
            "download_kbps": float(line[6]),
            "upload_kbps": float(line[7]),
            "total_tests": int(line[8]),
            "distance_miles": float(line[9])
        }
        counter = counter + 1
        data_str = json.dumps(record)
        data = data_str.encode("utf-8")
        publish_future = client.publish(topic, data)
        publish_future.add_done_callback(get_callback(publish_future, data))
        publish_futures.append(publish_future)
        print(counter)
        # sleep between 0 and 1 seconds
        sleep(random())

futures.wait(publish_futures, return_when=futures.ALL_COMPLETED)