1# Copyright 2014 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Script to populate datastore with system test data.""" 16 17 18from __future__ import print_function 19 20import os 21import string 22import sys 23import time 24import uuid 25 26from google.cloud import datastore 27 28 29ANCESTOR = ("Book", "GoT") 30RICKARD = ANCESTOR + ("Character", "Rickard") 31EDDARD = RICKARD + ("Character", "Eddard") 32KEY_PATHS = ( 33 RICKARD, 34 EDDARD, 35 ANCESTOR + ("Character", "Catelyn"), 36 EDDARD + ("Character", "Arya"), 37 EDDARD + ("Character", "Sansa"), 38 EDDARD + ("Character", "Robb"), 39 EDDARD + ("Character", "Bran"), 40 EDDARD + ("Character", "Jon Snow"), 41) 42CHARACTERS = ( 43 {"name": u"Rickard", "family": u"Stark", "appearances": 0, "alive": False}, 44 {"name": u"Eddard", "family": u"Stark", "appearances": 9, "alive": False}, 45 { 46 "name": u"Catelyn", 47 "family": [u"Stark", u"Tully"], 48 "appearances": 26, 49 "alive": False, 50 }, 51 {"name": u"Arya", "family": u"Stark", "appearances": 33, "alive": True}, 52 {"name": u"Sansa", "family": u"Stark", "appearances": 31, "alive": True}, 53 {"name": u"Robb", "family": u"Stark", "appearances": 22, "alive": False}, 54 {"name": u"Bran", "family": u"Stark", "appearances": 25, "alive": True}, 55 {"name": u"Jon Snow", "family": u"Stark", "appearances": 32, "alive": True}, 56) 57LARGE_CHARACTER_TOTAL_OBJECTS = 2500 58LARGE_CHARACTER_NAMESPACE = "LargeCharacterEntity" 59LARGE_CHARACTER_KIND = "LargeCharacter" 60 61 62def print_func(message): 63 if os.getenv("GOOGLE_CLOUD_NO_PRINT") != "true": 64 print(message) 65 66 67def add_large_character_entities(client=None): 68 MAX_STRING = (string.ascii_lowercase * 58)[:1500] 69 70 client.namespace = LARGE_CHARACTER_NAMESPACE 71 72 # Query used for all tests 73 page_query = client.query( 74 kind=LARGE_CHARACTER_KIND, namespace=LARGE_CHARACTER_NAMESPACE 75 ) 76 77 def put_objects(count): 78 current = 0 79 80 # Can only do 500 operations in a transaction with an overall 81 # size limit. 82 ENTITIES_TO_BATCH = 25 83 while current < count: 84 start = current 85 end = min(current + ENTITIES_TO_BATCH, count) 86 with client.transaction() as xact: 87 # The name/ID for the new entity 88 for i in range(start, end): 89 name = "character{0:05d}".format(i) 90 # The Cloud Datastore key for the new entity 91 task_key = client.key(LARGE_CHARACTER_KIND, name) 92 93 # Prepares the new entity 94 task = datastore.Entity(key=task_key) 95 task["name"] = "{0:05d}".format(i) 96 task["family"] = "Stark" 97 task["alive"] = False 98 99 for i in string.ascii_lowercase: 100 task["space-{}".format(i)] = MAX_STRING 101 102 # Saves the entity 103 xact.put(task) 104 current += ENTITIES_TO_BATCH 105 106 # Ensure we have 1500 entities for tests. If not, clean up type and add 107 # new entities equal to LARGE_CHARACTER_TOTAL_OBJECTS 108 all_entities = [e for e in page_query.fetch()] 109 if len(all_entities) != LARGE_CHARACTER_TOTAL_OBJECTS: 110 # Cleanup Collection if not an exact match 111 while all_entities: 112 entities = all_entities[:500] 113 all_entities = all_entities[500:] 114 client.delete_multi([e.key for e in entities]) 115 # Put objects 116 put_objects(LARGE_CHARACTER_TOTAL_OBJECTS) 117 118 119def add_characters(client=None): 120 if client is None: 121 # Get a client that uses the test dataset. 122 client = datastore.Client() 123 with client.transaction() as xact: 124 for key_path, character in zip(KEY_PATHS, CHARACTERS): 125 if key_path[-1] != character["name"]: 126 raise ValueError(("Character and key don't agree", key_path, character)) 127 entity = datastore.Entity(key=client.key(*key_path)) 128 entity.update(character) 129 xact.put(entity) 130 print_func( 131 "Adding Character %s %s" % (character["name"], character["family"]) 132 ) 133 134 135def add_uid_keys(client=None): 136 if client is None: 137 # Get a client that uses the test dataset. 138 client = datastore.Client() 139 140 num_batches = 2 141 batch_size = 500 142 143 for batch_num in range(num_batches): 144 with client.batch() as batch: 145 for seq_no in range(batch_size): 146 uid = str(uuid.uuid4()) 147 key = client.key("uuid_key", uid) 148 entity = datastore.Entity(key=key) 149 entity["batch_num"] = batch_num 150 entity["seq_no"] = seq_no 151 batch.put(entity) 152 153 154def add_timestamp_keys(client=None): 155 if client is None: 156 # Get a client that uses the test dataset. 157 client = datastore.Client() 158 159 num_batches = 20 160 batch_size = 500 161 162 timestamp_micros = set() 163 for batch_num in range(num_batches): 164 with client.batch() as batch: 165 for seq_no in range(batch_size): 166 print("time_time: batch: {}, sequence: {}".format(batch_num, seq_no)) 167 now_micros = int(time.time() * 1e6) 168 while now_micros in timestamp_micros: 169 now_micros = int(time.time() * 1e6) 170 timestamp_micros.add(now_micros) 171 key = client.key("timestamp_key", now_micros) 172 entity = datastore.Entity(key=key) 173 entity["batch_num"] = batch_num 174 entity["seq_no"] = seq_no 175 batch.put(entity) 176 177 178def main(): 179 client = datastore.Client() 180 flags = sys.argv[1:] 181 182 if len(flags) == 0: 183 flags = ["--characters", "--uuid", "--timestamps"] 184 185 if "--characters" in flags: 186 add_characters(client) 187 188 if "--uuid" in flags: 189 add_uid_keys(client) 190 191 if "--timestamps" in flags: 192 add_timestamp_keys(client) 193 194 195if __name__ == "__main__": 196 main() 197