File size: 5,259 Bytes
ccaf67e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Peewee migrations -- 002_add_local_sharing.py.

Some examples (model - class or model name)::

    > Model = migrator.orm['table_name']            # Return model in current state by name
    > Model = migrator.ModelClass                   # Return model in current state by name

    > migrator.sql(sql)                             # Run custom SQL
    > migrator.run(func, *args, **kwargs)           # Run python function with the given args
    > migrator.create_model(Model)                  # Create a model (could be used as decorator)
    > migrator.remove_model(model, cascade=True)    # Remove a model
    > migrator.add_fields(model, **fields)          # Add fields to a model
    > migrator.change_fields(model, **fields)       # Change fields
    > migrator.remove_fields(model, *field_names, cascade=True)
    > migrator.rename_field(model, old_field_name, new_field_name)
    > migrator.rename_table(model, new_table_name)
    > migrator.add_index(model, *col_names, unique=False)
    > migrator.add_not_null(model, *field_names)
    > migrator.add_default(model, field_name, default)
    > migrator.add_constraint(model, name, sql)
    > migrator.drop_index(model, *col_names)
    > migrator.drop_not_null(model, *field_names)
    > migrator.drop_constraints(model, *constraints)

"""

from contextlib import suppress

import peewee as pw
from peewee_migrate import Migrator


with suppress(ImportError):
    import playhouse.postgres_ext as pw_pext


def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
    """Write your migrations here."""

    if isinstance(database, pw.SqliteDatabase):
        migrate_sqlite(migrator, database, fake=fake)
    else:
        migrate_external(migrator, database, fake=fake)


def migrate_sqlite(migrator: Migrator, database: pw.Database, *, fake=False):
    # Adding fields created_at and updated_at to the 'chat' table
    migrator.add_fields(
        "chat",
        created_at=pw.DateTimeField(null=True),  # Allow null for transition
        updated_at=pw.DateTimeField(null=True),  # Allow null for transition
    )

    # Populate the new fields from an existing 'timestamp' field
    migrator.sql(
        "UPDATE chat SET created_at = timestamp, updated_at = timestamp WHERE timestamp IS NOT NULL"
    )

    # Now that the data has been copied, remove the original 'timestamp' field
    migrator.remove_fields("chat", "timestamp")

    # Update the fields to be not null now that they are populated
    migrator.change_fields(
        "chat",
        created_at=pw.DateTimeField(null=False),
        updated_at=pw.DateTimeField(null=False),
    )


def migrate_external(migrator: Migrator, database: pw.Database, *, fake=False):
    # Adding fields created_at and updated_at to the 'chat' table
    migrator.add_fields(
        "chat",
        created_at=pw.BigIntegerField(null=True),  # Allow null for transition
        updated_at=pw.BigIntegerField(null=True),  # Allow null for transition
    )

    # Populate the new fields from an existing 'timestamp' field
    migrator.sql(
        "UPDATE chat SET created_at = timestamp, updated_at = timestamp WHERE timestamp IS NOT NULL"
    )

    # Now that the data has been copied, remove the original 'timestamp' field
    migrator.remove_fields("chat", "timestamp")

    # Update the fields to be not null now that they are populated
    migrator.change_fields(
        "chat",
        created_at=pw.BigIntegerField(null=False),
        updated_at=pw.BigIntegerField(null=False),
    )


def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
    """Write your rollback migrations here."""

    if isinstance(database, pw.SqliteDatabase):
        rollback_sqlite(migrator, database, fake=fake)
    else:
        rollback_external(migrator, database, fake=fake)


def rollback_sqlite(migrator: Migrator, database: pw.Database, *, fake=False):
    # Recreate the timestamp field initially allowing null values for safe transition
    migrator.add_fields("chat", timestamp=pw.DateTimeField(null=True))

    # Copy the earliest created_at date back into the new timestamp field
    # This assumes created_at was originally a copy of timestamp
    migrator.sql("UPDATE chat SET timestamp = created_at")

    # Remove the created_at and updated_at fields
    migrator.remove_fields("chat", "created_at", "updated_at")

    # Finally, alter the timestamp field to not allow nulls if that was the original setting
    migrator.change_fields("chat", timestamp=pw.DateTimeField(null=False))


def rollback_external(migrator: Migrator, database: pw.Database, *, fake=False):
    # Recreate the timestamp field initially allowing null values for safe transition
    migrator.add_fields("chat", timestamp=pw.BigIntegerField(null=True))

    # Copy the earliest created_at date back into the new timestamp field
    # This assumes created_at was originally a copy of timestamp
    migrator.sql("UPDATE chat SET timestamp = created_at")

    # Remove the created_at and updated_at fields
    migrator.remove_fields("chat", "created_at", "updated_at")

    # Finally, alter the timestamp field to not allow nulls if that was the original setting
    migrator.change_fields("chat", timestamp=pw.BigIntegerField(null=False))