Restaurants JSON generator for MySQL Document Store

generate_restaurants.py

#!/usr/bin/python3.8

import mysqlx
import sys
from random import choice, randrange
from faker import Faker


def connect():
    session = mysqlx.get_session(
        {
            "host": "db instance IP",
            "port": 33060,
            "user": "login",
            "password": "password",
            "ssl-mode": "REQUIRED",
        }
    )
    return session


def gen_cuisine():
    cuisine_list = [
        "Belgian",
        "Italian"
    ]
    return choice(cuisine_list)


session = connect()
db = session.get_schema("docstore")
col = db.get_collection("restaurants")
fake = Faker()

print("Generating new documents.", end="", flush=True)

total = 1000

if len(sys.argv) > 1:
    if sys.argv[1]:
        total = int(sys.argv[1])

for _ in range(total):
    doc = {}
    doc["name"] = fake.company()
    address = {}
    address["street"] = fake.street_name()
    address["building"] = fake.building_number()
    address["zipcode"] = fake.postcode()
    doc["borough"] = fake.city()
    doc["cuisine"] = gen_cuisine()
    coord = []
    coord.append(float(fake.latitude()))
    coord.append(float(fake.longitude()))
    address["coord"] = coord
    doc["address"] = address
    grades = []
    for _ in range(randrange(5)):
        grade = {}
        grade_date = {}
        date = fake.date_time_this_decade()
        grade_date["$date"] = date.strftime("%Y-%m-%dT00:00:00.000+0000")
        grade_note = choice(["A", "B", "C"])
        grade_score = randrange(20)
        grade["date"] = grade_date
        grade["grade"] = grade_note
        grade["score"] = grade_score
        grades.append(grade)

    doc["grades"] = grades

    col.add(doc).execute()
    if total > 100000 and not _ % 1000:
        print(".", end="", flush=True)
    else:
        print(".", end="", flush=True)

print("\nDone ! {} documents generated.".format(total))