forked from googleapis/python-spanner-django
/
utils.py
136 lines (106 loc) · 4.09 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
# Copyright 2021 Google LLC
#
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file or at
# https://developers.google.com/open-source/licenses/bsd
import os
import time
from django.core.management import call_command
from django.db import connection
from google.api_core import exceptions
from google.cloud.spanner_v1 import Client
from google.cloud.spanner_v1.instance import Instance, Backup
from test_utils.retry import RetryErrors
from django_spanner.creation import DatabaseCreation
from django_spanner import USE_EMULATOR
CREATE_INSTANCE = (
os.getenv("GOOGLE_CLOUD_TESTS_CREATE_SPANNER_INSTANCE") is not None
)
SPANNER_OPERATION_TIMEOUT_IN_SECONDS = int(
os.getenv("SPANNER_OPERATION_TIMEOUT_IN_SECONDS", 60)
)
EXISTING_INSTANCES = []
INSTANCE_ID = os.environ.get(
"GOOGLE_CLOUD_TESTS_SPANNER_INSTANCE", "spanner-django-python-systest"
)
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "emulator-test-project")
DATABASE_NAME = os.getenv("DJANGO_SPANNER_DB", "django_test_db")
class Config(object):
"""Run-time configuration to be modified at set-up.
This is a mutable stand-in to allow test set-up to modify
global state.
"""
CLIENT = None
INSTANCE_CONFIG = None
INSTANCE = None
DATABASE = None
def _list_instances():
return list(Config.CLIENT.list_instances())
def setup_instance():
if USE_EMULATOR:
from google.auth.credentials import AnonymousCredentials
Config.CLIENT = Client(
project=PROJECT_ID, credentials=AnonymousCredentials()
)
else:
Config.CLIENT = Client()
retry = RetryErrors(exceptions.ServiceUnavailable)
configs = list(retry(Config.CLIENT.list_instance_configs)())
instances = retry(_list_instances)()
EXISTING_INSTANCES[:] = instances
# Delete test instances that are older than an hour.
cutoff = int(time.time()) - 1 * 60 * 60
instance_pbs = Config.CLIENT.list_instances(
"labels.python-spanner-systests:true"
)
for instance_pb in instance_pbs:
instance = Instance.from_pb(instance_pb, Config.CLIENT)
if "created" not in instance.labels:
continue
create_time = int(instance.labels["created"])
if create_time > cutoff:
continue
if not USE_EMULATOR:
# Instance cannot be deleted while backups exist.
for backup_pb in instance.list_backups():
backup = Backup.from_pb(backup_pb, instance)
backup.delete()
instance.delete()
if CREATE_INSTANCE:
if not USE_EMULATOR:
# Defend against back-end returning configs for regions we aren't
# actually allowed to use.
configs = [config for config in configs if "-us-" in config.name]
if not configs:
raise ValueError("List instance configs failed in module set up.")
Config.INSTANCE_CONFIG = configs[0]
config_name = configs[0].name
create_time = str(int(time.time()))
labels = {"django-spanner-systests": "true", "created": create_time}
Config.INSTANCE = Config.CLIENT.instance(
INSTANCE_ID, config_name, labels=labels
)
if not Config.INSTANCE.exists():
created_op = Config.INSTANCE.create()
created_op.result(
SPANNER_OPERATION_TIMEOUT_IN_SECONDS
) # block until completion
else:
Config.INSTANCE.reload()
else:
Config.INSTANCE = Config.CLIENT.instance(INSTANCE_ID)
Config.INSTANCE.reload()
def teardown_instance():
if CREATE_INSTANCE:
Config.INSTANCE.delete()
def setup_database():
Config.DATABASE = Config.INSTANCE.database(DATABASE_NAME)
if not Config.DATABASE.exists():
creation = DatabaseCreation(connection)
creation._create_test_db(verbosity=0, autoclobber=False, keepdb=True)
# Running migrations on the db.
call_command("migrate", interactive=False)
def teardown_database():
Config.DATABASE = Config.INSTANCE.database(DATABASE_NAME)
if Config.DATABASE.exists():
Config.DATABASE.drop()