Skip to content

Commit 5e29335

Browse files
tseavertseaver
committed
Scripts populating / cleaning DB for streaming / chunking systests.
Toward googleapis#3019.
1 parent ae98aca commit 5e29335

6 files changed

Lines changed: 187 additions & 1 deletion

File tree

‎spanner/nox.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ def system_tests(session, python_version):
6161
session.install('.')
6262

6363
# Run py.test against the system tests.
64-
session.run('py.test', '--quiet', 'tests/system.py')
64+
session.run('py.test', '--quiet', 'tests/system')
6565

6666

6767
@nox.session

‎spanner/tests/system/__init__.py‎

Whitespace-only changes.

spanner/tests/system.py renamed to spanner/tests/system/test_system.py

File renamed without changes.

‎spanner/tests/system/utils/__init__.py‎

Whitespace-only changes.

‎spanner/tests/system/utils/clear_streaming.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
# Copyright 2017 Google Inc.
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+
"""Depopulate spanner databases with data for streaming system tests."""
16+
17+
import os
18+
19+
from google.cloud.spanner import Client
20+
from google.cloud.spanner.keyset import KeySet
21+
from google.cloud.spanner.pool import BurstyPool
22+
23+
24+
INSTANCE_NAME = 'gcp-streaming-systests'
25+
DATABASE_NAME = 'testing'
26+
27+
28+
def print_func(message):
29+
if os.getenv('GOOGLE_CLOUD_NO_PRINT') != 'true':
30+
print(message)
31+
32+
33+
def remove_database(client):
34+
instance = client.instance(INSTANCE_NAME)
35+
36+
if not instance.exists():
37+
print_func("Instance does not exist: {}".format(INSTANCE_NAME))
38+
return
39+
40+
print_func("Instance exists: {}".format(INSTANCE_NAME))
41+
instance.reload()
42+
43+
pool = BurstyPool()
44+
database = instance.database(DATABASE_NAME)
45+
46+
if not database.exists():
47+
print_func("Database does not exist: {}".format(DATABASE_NAME))
48+
return
49+
print_func("Dropping database: {}".format(DATABASE_NAME))
50+
database.drop()
51+
52+
53+
if __name__ == '__main__':
54+
client = Client()
55+
remove_database(client)

‎spanner/tests/system/utils/populate_streaming.py‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
# Copyright 2017 Google Inc.
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+
"""Populate spanner databases with data for streaming system tests."""
16+
17+
import os
18+
19+
from google.cloud.spanner import Client
20+
from google.cloud.spanner.keyset import KeySet
21+
from google.cloud.spanner.pool import BurstyPool
22+
23+
24+
INSTANCE_NAME = 'gcp-streaming-systests'
25+
DATABASE_NAME = 'testing'
26+
27+
DDL = """\
28+
CREATE TABLE four_kay (
29+
pkey INT64,
30+
chunk_me STRING(4096) )
31+
PRIMARY KEY (pkey);
32+
CREATE TABLE forty_kay (
33+
pkey INT64,
34+
chunk_me STRING(40960) )
35+
PRIMARY KEY (pkey);
36+
CREATE TABLE four_hundred_kay (
37+
pkey INT64,
38+
chunk_me STRING(409600) )
39+
PRIMARY KEY (pkey);
40+
CREATE TABLE four_meg (
41+
pkey INT64,
42+
chunk_me STRING(2097152),
43+
chunk_me_2 STRING(2097152) )
44+
PRIMARY KEY (pkey);
45+
"""
46+
47+
DDL_STATEMENTS = [stmt.strip() for stmt in DDL.split(';') if stmt.strip()]
48+
49+
50+
def print_func(message):
51+
if os.getenv('GOOGLE_CLOUD_NO_PRINT') != 'true':
52+
print(message)
53+
54+
55+
def ensure_database(client):
56+
instance = client.instance(INSTANCE_NAME)
57+
58+
if not instance.exists():
59+
configs = list(client.list_instance_configs())
60+
config_name = configs[0].name
61+
print_func("Creating instance: {}".format(INSTANCE_NAME))
62+
instance = client.instance(INSTANCE_NAME, config_name)
63+
operation = instance.create()
64+
operation.result(30)
65+
else:
66+
print_func("Instance exists: {}".format(INSTANCE_NAME))
67+
instance.reload()
68+
69+
pool = BurstyPool()
70+
database = instance.database(
71+
DATABASE_NAME, ddl_statements=DDL_STATEMENTS, pool=pool)
72+
73+
if not database.exists():
74+
print_func("Creating database: {}".format(DATABASE_NAME))
75+
operation = database.create()
76+
operation.result(30)
77+
else:
78+
print_func("Database exists: {}".format(DATABASE_NAME))
79+
database.reload()
80+
81+
return database
82+
83+
84+
def populate_table(database, table_name, row_count, val_size):
85+
all_ = KeySet(all_=True)
86+
columns = ('pkey', 'chunk_me')
87+
rows = list(database.execute_sql(
88+
'SELECT COUNT(*) FROM {}'.format(table_name)))
89+
assert len(rows) == 1
90+
count = rows[0][0]
91+
if count != row_count:
92+
print_func("Repopulating table: {}".format(table_name))
93+
chunk_me = 'X' * val_size
94+
row_data = [(index, chunk_me) for index in range(row_count)]
95+
with database.batch() as batch:
96+
batch.delete(table_name, all_)
97+
batch.insert(table_name, columns, row_data)
98+
else:
99+
print_func("Leaving table: {}".format(table_name))
100+
101+
102+
def populate_table_2_columns(database, table_name, row_count, val_size):
103+
all_ = KeySet(all_=True)
104+
columns = ('pkey', 'chunk_me', 'chunk_me_2')
105+
rows = list(database.execute_sql(
106+
'SELECT COUNT(*) FROM {}'.format(table_name)))
107+
assert len(rows) == 1
108+
count = rows[0][0]
109+
if count != row_count:
110+
print_func("Repopulating table: {}".format(table_name))
111+
chunk_me = 'X' * val_size
112+
row_data = [(index, chunk_me, chunk_me) for index in range(row_count)]
113+
with database.batch() as batch:
114+
batch.delete(table_name, all_)
115+
batch.insert(table_name, columns, row_data)
116+
else:
117+
print_func("Leaving table: {}".format(table_name))
118+
119+
120+
def populate_streaming(client):
121+
database = ensure_database(client)
122+
populate_table(database, 'four_kay', 1000, 4096)
123+
populate_table(database, 'forty_kay', 100, 4096 * 10)
124+
populate_table(database, 'four_hundred_kay', 25, 4096 * 100)
125+
# Max STRING column size is just larger than 2 Mb, so use two columns
126+
populate_table_2_columns(database, 'four_meg', 10, 2048 * 1024)
127+
128+
129+
if __name__ == '__main__':
130+
client = Client()
131+
populate_streaming(client)

0 commit comments

Comments
 (0)