-
Notifications
You must be signed in to change notification settings - Fork 14
Expand file tree
/
Copy pathupdate.py
More file actions
361 lines (280 loc) · 12 KB
/
update.py
File metadata and controls
361 lines (280 loc) · 12 KB
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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Used with SpannerAdminApi to manage Spanner schema updates."""
import abc
from typing import Iterable, List, Optional, Type
from spanner_orm import condition
from spanner_orm import error
from spanner_orm import field
from spanner_orm import foreign_key_relationship
from spanner_orm import index
from spanner_orm import model
from spanner_orm.admin import api
from spanner_orm.admin import index_column
from spanner_orm.admin import metadata
class MigrationUpdate(abc.ABC):
"""Base class for all updates that can happen in a migration."""
@abc.abstractmethod
def execute(self) -> None:
"""Executes the update."""
raise NotImplementedError
class NoUpdate(MigrationUpdate):
"""Update that does nothing, for migrations that don't update db schemas."""
def execute(self) -> None:
"""See base class."""
class SchemaUpdate(MigrationUpdate, abc.ABC):
"""Base class for specifying schema updates."""
@abc.abstractmethod
def ddl(self) -> str:
raise NotImplementedError
def execute(self) -> None:
self.validate()
api.spanner_admin_api().update_schema(self.ddl())
def validate(self) -> None:
pass # TODO(dseomn): Remove this method.
class CreateTable(SchemaUpdate):
"""Update that allows creating a new table."""
def __init__(self, model_: Type[model.Model]):
self._model = model_
def ddl(self) -> str:
key_fields = [
'{} {}'.format(name, field.ddl())
for name, field in self._model.fields.items()
]
key_fields_ddl = ', '.join(key_fields)
for relation in self._model.foreign_key_relations.values():
key_fields_ddl += f', {relation.ddl}'
index_ddl = 'PRIMARY KEY ({})'.format(', '.join(self._model.primary_keys))
statement = 'CREATE TABLE {} ({}) {}'.format(self._model.table,
key_fields_ddl, index_ddl)
if self._model.interleaved:
statement += ', INTERLEAVE IN PARENT {parent} ON DELETE CASCADE'.format(
parent=self._model.interleaved.table)
return statement
def validate(self) -> None:
if not self._model.table:
raise error.SpannerError('New table has no name')
existing_model = metadata.SpannerMetadata.model(self._model.table)
if existing_model:
raise error.SpannerError('Table {} already exists'.format(
self._model.table))
if self._model.interleaved:
self._validate_parent()
self._validate_primary_keys()
if self._model.indexes.keys() - {index.Index.PRIMARY_INDEX}:
raise error.SpannerError(
'Secondary indexes cannot be created by CreateTable; use CreateIndex '
'in a separate migration.')
def _validate_parent(self) -> None:
"""Verifies that the parent table information is valid."""
parent_primary_keys = self._model.interleaved.primary_keys
primary_keys = self._model.primary_keys
message = 'Table {} is not a child of parent table {}'.format(
self._model.table, self._model.interleaved.table)
for parent_key, key in zip(parent_primary_keys, primary_keys):
if parent_key != key:
raise error.SpannerError(message)
if len(parent_primary_keys) > len(primary_keys):
raise error.SpannerError(message)
def _validate_primary_keys(self) -> None:
"""Verifies that the primary key data is valid."""
if not self._model.primary_keys:
raise error.SpannerError('Table {} has no primary key'.format(
self._model.table))
for key in self._model.primary_keys:
if key not in self._model.fields:
raise error.SpannerError(
'Table {} column {} in primary key but not in schema'.format(
self._model.table, key))
class DropTable(SchemaUpdate):
"""Update for dropping an existing table."""
def __init__(self, table_name: str):
self._table = table_name
def ddl(self) -> str:
return 'DROP TABLE {}'.format(self._table)
class AddColumn(SchemaUpdate):
"""Update for adding a column to an existing table.
Only supports adding nullable columns
"""
def __init__(self, table_name: str, column_name: str, field_: field.Field):
self._table = table_name
self._column = column_name
self._field = field_
def ddl(self) -> str:
return 'ALTER TABLE {} ADD COLUMN {} {}'.format(self._table, self._column,
self._field.ddl())
def validate(self) -> None:
model_ = metadata.SpannerMetadata.model(self._table)
if not model_:
raise error.SpannerError('Table {} does not exist'.format(self._table))
if not self._field.nullable():
raise error.SpannerError('Column {} is not nullable'.format(self._column))
if self._field.primary_key():
raise error.SpannerError('Column {} is a primary key'.format(
self._column))
class DropColumn(SchemaUpdate):
"""Update for dropping a column from an existing table."""
def __init__(self, table_name: str, column_name: str):
self._table = table_name
self._column = column_name
def ddl(self) -> str:
return 'ALTER TABLE {} DROP COLUMN {}'.format(self._table, self._column)
def validate(self) -> None:
model_ = metadata.SpannerMetadata.model(self._table)
if not model_:
raise error.SpannerError('Table {} does not exist'.format(self._table))
if self._column not in model_.fields:
raise error.SpannerError('Column {} does not exist on {}'.format(
self._column, self._table))
# Verify no indices exist on the column we're trying to drop
num_indexed_columns = index_column.IndexColumnSchema.count(
condition.equal_to('column_name', self._column),
condition.equal_to('table_name', self._table),
)
if num_indexed_columns > 0:
raise error.SpannerError('Column {} is indexed'.format(self._column))
class AlterColumn(SchemaUpdate):
"""Update for altering a column an existing table.
Only supports changing the nullability of a column
"""
def __init__(self, table_name: str, column_name: str, field_: field.Field):
self._table = table_name
self._column = column_name
self._field = field_
def ddl(self) -> str:
return 'ALTER TABLE {} ALTER COLUMN {} {}'.format(self._table, self._column,
self._field.ddl())
def validate(self) -> None:
model_ = metadata.SpannerMetadata.model(self._table)
if not model_:
raise error.SpannerError('Table {} does not exist'.format(self._table))
if self._column not in model_.fields:
raise error.SpannerError('Column {} does not exist on {}'.format(
self._column, self._table))
if self._column in model_.primary_keys:
raise error.SpannerError('Column {} is a primary key on {}'.format(
self._column, self._table))
old_field = model_.fields[self._column]
# Validate that the only alteration is to change column nullability
if self._field.field_type().ddl() != old_field.field_type().ddl():
raise error.SpannerError('Column {} is changing type'.format(
self._column))
if self._field.nullable() == old_field.nullable():
raise error.SpannerError('Column {} has no changes'.format(self._column))
class CreateIndex(SchemaUpdate):
"""Update for creating an index on an existing table."""
def __init__(self,
table_name: str,
index_name: str,
columns: Iterable[str],
interleaved: Optional[str] = None,
null_filtered: bool = False,
unique: bool = False,
storing_columns: Optional[Iterable[str]] = None):
self._table = table_name
self._index = index_name
self._columns = columns
self._parent_table = interleaved
self._null_filtered = null_filtered
self._unique = unique
self._storing_columns = storing_columns or []
def ddl(self) -> str:
statement = 'CREATE'
if self._unique:
statement += ' UNIQUE'
if self._null_filtered:
statement += ' NULL_FILTERED'
statement += (f' INDEX {self._index} '
f'ON {self._table} ({", ".join(self._columns)})')
if self._storing_columns:
statement += 'STORING ({})'.format(', '.join(self._storing_columns))
if self._parent_table:
statement += ', INTERLEAVE IN {}'.format(self._parent_table)
return statement
def validate(self) -> None:
model_ = metadata.SpannerMetadata.model(self._table)
if not model_:
raise error.SpannerError('Table {} does not exist'.format(self._table))
if not self._columns:
raise error.SpannerError('Index {} has no columns'.format(self._index))
if self._index in model_.indexes:
raise error.SpannerError('Index {} already exists'.format(self._index))
self._validate_columns(model_)
if self._parent_table:
self._validate_parent(model_)
def _validate_columns(self, model_: Type[model.Model]) -> None:
"""Verifies all columns exist and are not part of the primary key."""
for column in self._columns:
if column not in model_.columns:
raise error.SpannerError('Table {} has no column {}'.format(
self._table, column))
for column in self._storing_columns:
if column not in model_.columns:
raise error.SpannerError('Table {} has no column {}'.format(
self._table, column))
if column in model_.primary_keys:
raise error.SpannerError('{} is part of the primary key for {}'.format(
column, self._table))
def _validate_parent(self, model_: Type[model.Model]) -> None:
"""Verifies this index can be interleaved in the parent table."""
parent = model_.interleaved
while parent:
if parent == self._parent_table:
break
parent = parent.interleaved
if not parent:
raise error.SpannerError('{} is not a parent of table {}'.format(
self._parent_table, self._table))
class DropIndex(SchemaUpdate):
"""Update for dropping a secondary index on an existing table."""
def __init__(self, table_name: str, index_name: str):
self._table = table_name
self._index = index_name
def ddl(self) -> str:
return 'DROP INDEX {}'.format(self._index)
def validate(self) -> None:
model_ = metadata.SpannerMetadata.model(self._table)
if not model_:
raise error.SpannerError('Table {} does not exist'.format(self._table))
db_index = model_.indexes.get(self._index)
if not db_index:
raise error.SpannerError('Index {} does not exist'.format(self._index))
if db_index.primary:
raise error.SpannerError('Index {} is the primary index'.format(
self._index))
class ExecutePartitionedDml(MigrationUpdate):
"""Update for running arbitrary partitioned DML.
NOTE: Partitioned DML queries should be idempotent. See
https://cloud.google.com/spanner/docs/dml-partitioned for details, and more
information about partitioned DML.
"""
def __init__(self, dml: str):
self._dml = dml
def execute(self) -> None:
"""See base class."""
api.spanner_admin_api().execute_partitioned_dml(self._dml)
def model_creation_ddl(model_: Type[model.Model]) -> List[str]:
"""Returns the list of ddl statements needed to create the model's table."""
ddl_list = [CreateTable(model_).ddl()]
for model_index in model_.indexes.values():
if model_index.primary:
continue
create_index = CreateIndex(
model_.table,
model_index.name,
model_index.columns,
interleaved=model_index.parent,
storing_columns=model_index.storing_columns)
ddl_list.append(create_index.ddl())
return ddl_list