client SDK performance tester for Cloud Datastore

datastore-test.py

import datetime
import time
import random
from google.cloud import datastore

def _rand() -> int:
    """Returns a random integer centered around 0"""
    return random.randint(0, 6) - 3

prefix: str = f'[python-datastore@{datastore.__version__}]'
unique_key = time.time()
batch_id: str = "batch-{unique_key}"

def log_time(start: float, finish: float, msg: str):
    runtime = finish - start
    print(f'{prefix} finished {msg} in {round(runtime, 3) * 1000} ms')

client = datastore.Client()
chunk = []
keys = []

loop_st = time.time()
counter = 0
while counter < 500:
    customer_key = client.key('Customer', f'Batch-{unique_key}-{counter}')
    keys.append(customer_key)
    customer = datastore.Entity(customer_key)
    customer.update({
        "batch": batch_id,
        "tags": ["fun", "programming"],
        "collaborators": ["alice", "bob"],
        "signed_up_at": datetime.datetime.now() - datetime.timedelta(seconds=counter + _rand()),
    })
    chunk.append(customer)
    counter += 1
log_time(loop_st, time.time(), 'creation loop')

st = time.time()
with client.batch():
    client.put_multi(chunk)
log_time(st, time.time(), 'batch write')

entities_for_update = []

st = time.time()
with client.batch():
    entities_for_update = client.get_multi(keys)
log_time(st, time.time(), 'batch get')

for entity in entities_for_update:
    entity.update({'is_updated': True})

st = time.time()
with client.batch():
    client.put_multi(entities_for_update)
log_time(st, time.time(), 'batch update')

st = time.time()
query = client.query(kind="Customer")
query = query.add_filter("batch", "=", batch_id)
query.order = ["signed_up_at"]
records = list(query.fetch())
log_time(st, time.time(), 'batch query')

print(f'  > loaded {len(records)} objects')