008_add_models.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. """Peewee migrations -- 002_add_local_sharing.py.
  2. Some examples (model - class or model name)::
  3. > Model = migrator.orm['table_name'] # Return model in current state by name
  4. > Model = migrator.ModelClass # Return model in current state by name
  5. > migrator.sql(sql) # Run custom SQL
  6. > migrator.run(func, *args, **kwargs) # Run python function with the given args
  7. > migrator.create_model(Model) # Create a model (could be used as decorator)
  8. > migrator.remove_model(model, cascade=True) # Remove a model
  9. > migrator.add_fields(model, **fields) # Add fields to a model
  10. > migrator.change_fields(model, **fields) # Change fields
  11. > migrator.remove_fields(model, *field_names, cascade=True)
  12. > migrator.rename_field(model, old_field_name, new_field_name)
  13. > migrator.rename_table(model, new_table_name)
  14. > migrator.add_index(model, *col_names, unique=False)
  15. > migrator.add_not_null(model, *field_names)
  16. > migrator.add_default(model, field_name, default)
  17. > migrator.add_constraint(model, name, sql)
  18. > migrator.drop_index(model, *col_names)
  19. > migrator.drop_not_null(model, *field_names)
  20. > migrator.drop_constraints(model, *constraints)
  21. """
  22. from contextlib import suppress
  23. import peewee as pw
  24. from peewee_migrate import Migrator
  25. with suppress(ImportError):
  26. import playhouse.postgres_ext as pw_pext
  27. def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
  28. """Write your migrations here."""
  29. # Adding fields created_at and updated_at to the 'user' table
  30. migrator.add_fields(
  31. "user",
  32. created_at=pw.BigIntegerField(null=True), # Allow null for transition
  33. updated_at=pw.BigIntegerField(null=True), # Allow null for transition
  34. last_active_at=pw.BigIntegerField(null=True), # Allow null for transition
  35. )
  36. # Populate the new fields from an existing 'timestamp' field
  37. migrator.sql(
  38. 'UPDATE "user" SET created_at = timestamp, updated_at = timestamp, last_active_at = timestamp WHERE timestamp IS NOT NULL'
  39. )
  40. # Now that the data has been copied, remove the original 'timestamp' field
  41. migrator.remove_fields("user", "timestamp")
  42. # Update the fields to be not null now that they are populated
  43. migrator.change_fields(
  44. "user",
  45. created_at=pw.BigIntegerField(null=False),
  46. updated_at=pw.BigIntegerField(null=False),
  47. last_active_at=pw.BigIntegerField(null=False),
  48. )
  49. def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
  50. """Write your rollback migrations here."""
  51. # Recreate the timestamp field initially allowing null values for safe transition
  52. migrator.add_fields("user", timestamp=pw.BigIntegerField(null=True))
  53. # Copy the earliest created_at date back into the new timestamp field
  54. # This assumes created_at was originally a copy of timestamp
  55. migrator.sql('UPDATE "user" SET timestamp = created_at')
  56. # Remove the created_at and updated_at fields
  57. migrator.remove_fields("user", "created_at", "updated_at", "last_active_at")
  58. # Finally, alter the timestamp field to not allow nulls if that was the original setting
  59. migrator.change_fields("user", timestamp=pw.BigIntegerField(null=False))