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