Skip to content

Commit ce999f4

Browse files
dantownsenddantownsend
committed
can reverse alter statements in auto migrations
1 parent 4df6d4b commit ce999f4

File tree

4 files changed

+151
-189
lines changed

4 files changed

+151
-189
lines changed

‎piccolo/apps/migrations/auto/migration_manager.py‎

Lines changed: 148 additions & 188 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,7 @@ def alter_column(
233233
tablename: str,
234234
column_name: str,
235235
params: t.Dict[str, t.Any],
236+
old_params: t.Dict[str, t.Any],
236237
):
237238
"""
238239
All possible alterations aren't currently supported.
@@ -243,6 +244,7 @@ def alter_column(
243244
tablename=tablename,
244245
column_name=column_name,
245246
params=params,
247+
old_params=old_params,
246248
)
247249
)
248250

@@ -321,25 +323,123 @@ def deserialise_params(
321323

322324
###########################################################################
323325

324-
async def run(self):
325-
print("Running MigrationManager ...")
326+
async def _run_alter_columns(self, backwards=False):
327+
for table_class_name in self.alter_columns.table_class_names:
328+
alter_columns = self.alter_columns.for_table_class_name(
329+
table_class_name
330+
)
326331

327-
engine = engine_finder()
332+
if not alter_columns:
333+
continue
328334

329-
if not engine:
330-
raise Exception("Can't find engine")
335+
_Table: t.Type[Table] = type(table_class_name, (Table,), {})
336+
_Table._meta.tablename = alter_columns[0].tablename
331337

332-
async with engine.transaction():
338+
for column in alter_columns:
339+
params = column.old_params if backwards else column.params
340+
row_name = column.row_name
333341

334-
for raw in self.raw:
335-
if inspect.iscoroutinefunction(raw):
336-
await raw()
337-
else:
338-
raw()
342+
null = params.get("null")
343+
if null is not None:
344+
await _Table.alter().set_null(
345+
column=row_name, boolean=null
346+
).run()
347+
348+
length = params.get("length")
349+
if length is not None:
350+
await _Table.alter().set_length(
351+
column=row_name, length=length
352+
).run()
353+
354+
unique = params.get("unique")
355+
if unique is not None:
356+
await _Table.alter().set_unique(
357+
column=row_name, boolean=unique
358+
).run()
359+
360+
async def _run_drop_tables(self, backwards=False):
361+
if backwards:
362+
print("Dropped tables can't currently be reversed.")
363+
else:
364+
for table in self.drop_tables:
365+
await table.to_table_class().alter().drop_table().run()
366+
367+
async def _run_drop_columns(self, backwards=False):
368+
if backwards:
369+
print("Dropped columns can't currently be reversed.")
370+
else:
371+
for table_class_name in self.drop_columns.table_class_names:
372+
columns = self.drop_columns.for_table_class_name(
373+
table_class_name
374+
)
375+
376+
if not columns:
377+
continue
339378

340-
###################################################################
341-
# Add tables
379+
_Table: t.Type[Table] = type(table_class_name, (Table,), {})
380+
_Table._meta.tablename = columns[0].tablename
342381

382+
for column in columns:
383+
await _Table.alter().drop_column(
384+
column=column.column_name
385+
).run()
386+
387+
async def _run_rename_tables(self, backwards=False):
388+
for rename_table in self.rename_tables:
389+
class_name = (
390+
rename_table.new_class_name
391+
if backwards
392+
else rename_table.old_class_name
393+
)
394+
tablename = (
395+
rename_table.new_tablename
396+
if backwards
397+
else rename_table.old_tablename
398+
)
399+
new_tablename = (
400+
rename_table.old_tablename
401+
if backwards
402+
else rename_table.new_tablename
403+
)
404+
405+
_Table: t.Type[Table] = type(class_name, (Table,), {})
406+
_Table._meta.tablename = tablename
407+
408+
await _Table.alter().rename_table(new_name=new_tablename).run()
409+
410+
async def _run_rename_columns(self, backwards=False):
411+
for table_class_name in self.rename_columns.table_class_names:
412+
columns = self.rename_columns.for_table_class_name(
413+
table_class_name
414+
)
415+
416+
if not columns:
417+
continue
418+
419+
_Table: t.Type[Table] = type(table_class_name, (Table,), {})
420+
_Table._meta.tablename = columns[0].tablename
421+
422+
for rename_column in columns:
423+
column = (
424+
rename_column.new_column_name
425+
if backwards
426+
else rename_column.old_column_name
427+
)
428+
new_name = (
429+
rename_column.old_column_name
430+
if backwards
431+
else rename_column.new_column_name
432+
)
433+
434+
await _Table.alter().rename_column(
435+
column=column, new_name=new_name,
436+
).run()
437+
438+
async def _run_add_tables(self, backwards=False):
439+
if backwards:
440+
for add_table in self.add_tables:
441+
await add_table.to_table_class().alter().drop_table().run()
442+
else:
343443
for add_table in self.add_tables:
344444
columns = (
345445
self.add_columns.columns_for_table_class_name(
@@ -356,27 +456,19 @@ async def run(self):
356456

357457
await _Table.create_table().run()
358458

359-
###################################################################
360-
# Drop tables
361-
362-
for table in self.drop_tables:
363-
await table.to_table_class().alter().drop_table().run()
364-
365-
###################################################################
366-
# Rename tables
367-
368-
for rename_table in self.rename_tables:
459+
async def _run_add_columns(self, backwards=False):
460+
"""
461+
Add columns, which belong to existing tables
462+
"""
463+
if backwards:
464+
for add_column in self.add_columns.add_columns:
369465
_Table: t.Type[Table] = type(
370-
rename_table.old_class_name, (Table,), {}
466+
add_column.table_class_name, (Table,), {}
371467
)
372-
_Table._meta.tablename = rename_table.old_tablename
373-
await _Table.alter().rename_table(
374-
new_name=rename_table.new_tablename
375-
).run()
376-
377-
###################################################################
378-
# Add columns, which belong to existing tables
468+
_Table._meta.tablename = add_column.tablename
379469

470+
await _Table.alter().drop_column(add_column.column).run()
471+
else:
380472
new_table_class_names = [i.class_name for i in self.add_tables]
381473

382474
for table_class_name in self.add_columns.table_class_names:
@@ -398,82 +490,29 @@ async def run(self):
398490
name=column._meta.name, column=column
399491
).run()
400492

401-
###################################################################
402-
# Drop columns
403-
404-
for table_class_name in self.drop_columns.table_class_names:
405-
columns = self.drop_columns.for_table_class_name(
406-
table_class_name
407-
)
408-
409-
if not columns:
410-
continue
411-
412-
_Table: t.Type[Table] = type(table_class_name, (Table,), {})
413-
_Table._meta.tablename = columns[0].tablename
414-
415-
for column in columns:
416-
await _Table.alter().drop_column(
417-
column=column.column_name
418-
).run()
419-
420-
###################################################################
421-
# Rename columns
422-
423-
for table_class_name in self.rename_columns.table_class_names:
424-
columns = self.rename_columns.for_table_class_name(
425-
table_class_name
426-
)
427-
428-
if not columns:
429-
continue
430-
431-
_Table: t.Type[Table] = type(table_class_name, (Table,), {})
432-
_Table._meta.tablename = columns[0].tablename
433-
434-
for column in columns:
435-
await _Table.alter().rename_column(
436-
column=column.old_column_name,
437-
new_name=column.new_column_name,
438-
).run()
493+
async def run(self):
494+
print("Running MigrationManager ...")
439495

440-
###################################################################
441-
# Alter columns
496+
engine = engine_finder()
442497

443-
for table_class_name in self.alter_columns.table_class_names:
444-
alter_columns = self.alter_columns.for_table_class_name(
445-
table_class_name
446-
)
498+
if not engine:
499+
raise Exception("Can't find engine")
447500

448-
if not alter_columns:
449-
continue
501+
async with engine.transaction():
450502

451-
_Table: t.Type[Table] = type(table_class_name, (Table,), {})
452-
_Table._meta.tablename = alter_columns[0].tablename
453-
454-
for column in alter_columns:
455-
params = column.params
456-
row_name = column.row_name
457-
458-
null = params.get("null")
459-
if null is not None:
460-
await _Table.alter().set_null(
461-
column=row_name, boolean=null
462-
).run()
463-
464-
length = params.get("length")
465-
if length is not None:
466-
await _Table.alter().set_length(
467-
column=row_name, length=length
468-
).run()
469-
470-
unique = params.get("unique")
471-
if unique is not None:
472-
await _Table.alter().set_unique(
473-
column=row_name, boolean=unique
474-
).run()
503+
for raw in self.raw:
504+
if inspect.iscoroutinefunction(raw):
505+
await raw()
506+
else:
507+
raw()
475508

476-
###########################################################################
509+
await self._run_add_tables()
510+
await self._run_add_columns()
511+
await self._run_drop_tables()
512+
await self._run_rename_tables()
513+
await self._run_drop_columns()
514+
await self._run_rename_columns()
515+
await self._run_alter_columns()
477516

478517
async def run_backwards(self):
479518
print("Reversing MigrationManager ...")
@@ -491,89 +530,10 @@ async def run_backwards(self):
491530
else:
492531
raw()
493532

494-
###################################################################
495-
# Reverse add tables
496-
497-
for add_table in self.add_tables:
498-
await add_table.to_table_class().alter().drop_table().run()
499-
500-
###################################################################
501-
# Reverse drop tables
502-
503-
if self.drop_tables:
504-
print("Dropped tables can't currently be reversed.")
505-
506-
###################################################################
507-
# Reverse rename tables
508-
509-
for rename_table in self.rename_tables:
510-
_Table: t.Type[Table] = type(
511-
rename_table.new_class_name, (Table,), {}
512-
)
513-
_Table._meta.tablename = rename_table.new_tablename
514-
515-
await _Table.alter().rename_table(
516-
new_name=rename_table.old_tablename
517-
).run()
518-
519-
###################################################################
520-
# Reverse add columns
521-
522-
for add_column in self.add_columns.add_columns:
523-
_Table: t.Type[Table] = type(
524-
add_column.table_class_name, (Table,), {}
525-
)
526-
_Table._meta.tablename = add_column.tablename
527-
528-
await _Table.alter().drop_column(add_column.column).run()
529-
530-
###################################################################
531-
# Reverse drop columns
532-
533-
if self.drop_columns:
534-
print("Dropped columns can't currently be reversed.")
535-
536-
###################################################################
537-
# Reverse rename columns
538-
539-
for rename_column in self.rename_columns.rename_columns:
540-
_Table: t.Type[Table] = type(
541-
rename_column.table_class_name, (Table,), {}
542-
)
543-
_Table._meta.tablename = rename_column.tablename
544-
545-
await _Table.alter().rename_column(
546-
column=rename_column.new_column_name,
547-
new_name=rename_column.old_column_name,
548-
).run()
549-
550-
###################################################################
551-
# Alter columns
552-
553-
# TODO - need to find what the old values are.
554-
# for alter_column in self.alter_columns.alter_columns:
555-
# _Table: t.Type[Table] = type(
556-
# alter_column.table_class_name, (Table,), {}
557-
# )
558-
# _Table._meta.tablename = alter_column.tablename
559-
560-
# params = alter_column.params
561-
# row_name = alter_column.row_name
562-
563-
# null = params.get("null")
564-
# if null is not None:
565-
# await _Table.alter().set_null(
566-
# column=row_name, boolean=null
567-
# ).run()
568-
569-
# length = params.get("length")
570-
# if length is not None:
571-
# await _Table.alter().set_length(
572-
# column=row_name, length=length
573-
# ).run()
574-
575-
# unique = params.get("unique")
576-
# if unique is not None:
577-
# await _Table.alter().set_unique(
578-
# column=row_name, boolean=unique
579-
# ).run()
533+
await self._run_add_tables(backwards=True)
534+
await self._run_add_columns(backwards=True)
535+
await self._run_drop_tables(backwards=True)
536+
await self._run_rename_tables(backwards=True)
537+
await self._run_drop_columns(backwards=True)
538+
await self._run_rename_columns(backwards=True)
539+
await self._run_alter_columns(backwards=True)

‎piccolo/apps/migrations/auto/operations.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ class AlterColumn:
2424
column_name: str
2525
tablename: str
2626
params: t.Dict[str, t.Any]
27+
old_params: t.Dict[str, t.Any]
2728

2829

2930
@dataclass

0 commit comments

Comments
 (0)