gitscope
/
← marzneshin / app/db/migrations/versions/20240627_a83e4dd22672_refactor_user_expire.py / blame
Files Commits 446 Branches 1 Tags 43 Markers 4 README License
https://dev.kronoz.win/r/marzneshin.git
app/db/migrations/versions/20240627_a83e4dd22672_refactor_user_expire.py · nightly
Source
1 d377c4c Dawsh 2y ago """refactor user expire
2 d377c4c Dawsh 2y ago
3 d377c4c Dawsh 2y ago Revision ID: a83e4dd22672
4 d377c4c Dawsh 2y ago Revises: 9001356f0063
5 d377c4c Dawsh 2y ago Create Date: 2024-06-27 01:04:34.974520
6 d377c4c Dawsh 2y ago
7 d377c4c Dawsh 2y ago """
8 d377c4c Dawsh 2y ago
9 d377c4c Dawsh 2y ago import sqlalchemy as sa
10 d377c4c Dawsh 2y ago from alembic import op
11 d377c4c Dawsh 2y ago from sqlalchemy import table, column
12 d377c4c Dawsh 2y ago from sqlalchemy.dialects import postgresql
13 d377c4c Dawsh 2y ago
14 d377c4c Dawsh 2y ago # revision identifiers, used by Alembic.
15 d377c4c Dawsh 2y ago revision = "a83e4dd22672"
16 d377c4c Dawsh 2y ago down_revision = "9001356f0063"
17 d377c4c Dawsh 2y ago branch_labels = None
18 d377c4c Dawsh 2y ago depends_on = None
19 d377c4c Dawsh 2y ago
20 d377c4c Dawsh 2y ago
21 d377c4c Dawsh 2y ago def upgrade() -> None:
22 d377c4c Dawsh 2y ago bind = op.get_bind()
23 d377c4c Dawsh 2y ago dialect = bind.dialect.name
24 d377c4c Dawsh 2y ago
25 d377c4c Dawsh 2y ago if dialect == "postgresql":
26 d377c4c Dawsh 2y ago op.execute(
27 d377c4c Dawsh 2y ago """
28 d377c4c Dawsh 2y ago DO $$
29 d377c4c Dawsh 2y ago BEGIN
30 d377c4c Dawsh 2y ago IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'userexpirestrategy') THEN
31 d377c4c Dawsh 2y ago CREATE TYPE userexpirestrategy AS ENUM ('NEVER', 'FIXED_DATE', 'START_ON_FIRST_USE');
32 d377c4c Dawsh 2y ago END IF;
33 d377c4c Dawsh 2y ago END $$;
34 d377c4c Dawsh 2y ago """
35 d377c4c Dawsh 2y ago )
36 d377c4c Dawsh 2y ago op.add_column(
37 d377c4c Dawsh 2y ago "users",
38 d377c4c Dawsh 2y ago sa.Column(
39 d377c4c Dawsh 2y ago "activated",
40 d377c4c Dawsh 2y ago sa.Boolean(),
41 d377c4c Dawsh 2y ago nullable=False,
42 d377c4c Dawsh 2y ago server_default=sa.sql.true(),
43 d377c4c Dawsh 2y ago ),
44 d377c4c Dawsh 2y ago )
45 d377c4c Dawsh 2y ago op.add_column(
46 d377c4c Dawsh 2y ago "users",
47 d377c4c Dawsh 2y ago sa.Column(
48 d377c4c Dawsh 2y ago "expire_strategy",
49 d377c4c Dawsh 2y ago sa.Enum(
50 d377c4c Dawsh 2y ago "NEVER",
51 d377c4c Dawsh 2y ago "FIXED_DATE",
52 d377c4c Dawsh 2y ago "START_ON_FIRST_USE",
53 d377c4c Dawsh 2y ago name="userexpirestrategy",
54 d377c4c Dawsh 2y ago ),
55 d377c4c Dawsh 2y ago server_default="NEVER",
56 d377c4c Dawsh 2y ago nullable=False,
57 d377c4c Dawsh 2y ago ),
58 d377c4c Dawsh 2y ago )
59 d377c4c Dawsh 2y ago
60 d377c4c Dawsh 2y ago if dialect == "sqlite":
61 d377c4c Dawsh 2y ago for old, new, type_ in [
62 d377c4c Dawsh 2y ago ("on_hold_timeout", "activation_deadline", sa.DateTime),
63 d377c4c Dawsh 2y ago ("expire", "expire_date", sa.DateTime),
64 d377c4c Dawsh 2y ago ("on_hold_expire_duration", "usage_duration", sa.Integer),
65 d377c4c Dawsh 2y ago ]:
66 d377c4c Dawsh 2y ago op.add_column("users", sa.Column(new, type_))
67 d377c4c Dawsh 2y ago
68 d377c4c Dawsh 2y ago # Step 2: Copy data from the old column to the new column
69 d377c4c Dawsh 2y ago op.execute(f"UPDATE users SET {new} = {old}")
70 d377c4c Dawsh 2y ago
71 d377c4c Dawsh 2y ago with op.batch_alter_table("users") as batch_op:
72 d377c4c Dawsh 2y ago batch_op.drop_column(old)
73 d377c4c Dawsh 2y ago else:
74 d377c4c Dawsh 2y ago op.alter_column(
75 deb4d55 Dawsh 2y ago "users",
76 deb4d55 Dawsh 2y ago "on_hold_timeout",
77 deb4d55 Dawsh 2y ago new_column_name="activation_deadline",
78 deb4d55 Dawsh 2y ago existing_type=sa.DateTime,
79 deb4d55 Dawsh 2y ago )
80 deb4d55 Dawsh 2y ago op.alter_column(
81 deb4d55 Dawsh 2y ago "users",
82 deb4d55 Dawsh 2y ago "expire",
83 deb4d55 Dawsh 2y ago new_column_name="expire_date",
84 deb4d55 Dawsh 2y ago existing_type=sa.DateTime,
85 d377c4c Dawsh 2y ago )
86 d377c4c Dawsh 2y ago op.alter_column(
87 d377c4c Dawsh 2y ago "users",
88 d377c4c Dawsh 2y ago "on_hold_expire_duration",
89 d377c4c Dawsh 2y ago new_column_name="usage_duration",
90 deb4d55 Dawsh 2y ago existing_type=sa.BigInteger,
91 d377c4c Dawsh 2y ago )
92 d377c4c Dawsh 2y ago # ### end Alembic commands ###
93 d377c4c Dawsh 2y ago users_table = table(
94 d377c4c Dawsh 2y ago "users",
95 d377c4c Dawsh 2y ago column("id", sa.Integer),
96 d377c4c Dawsh 2y ago column("status", sa.String),
97 d377c4c Dawsh 2y ago column("expire_date", sa.DateTime),
98 d377c4c Dawsh 2y ago column("expire_strategy", sa.Enum(name="userexpirestrategy")),
99 d377c4c Dawsh 2y ago )
100 d377c4c Dawsh 2y ago
101 d377c4c Dawsh 2y ago connection = op.get_bind()
102 d377c4c Dawsh 2y ago results = connection.execute(
103 d377c4c Dawsh 2y ago sa.select(
104 d377c4c Dawsh 2y ago users_table.c.id, users_table.c.status, users_table.c.expire_date
105 d377c4c Dawsh 2y ago )
106 d377c4c Dawsh 2y ago )
107 d377c4c Dawsh 2y ago
108 d377c4c Dawsh 2y ago for user in results:
109 d377c4c Dawsh 2y ago status = user[1]
110 d377c4c Dawsh 2y ago expire_date = user[2]
111 d377c4c Dawsh 2y ago
112 d377c4c Dawsh 2y ago if status == "on_hold":
113 d377c4c Dawsh 2y ago expire_strategy = "START_ON_FIRST_USE"
114 d377c4c Dawsh 2y ago elif expire_date:
115 d377c4c Dawsh 2y ago expire_strategy = "FIXED_DATE"
116 d377c4c Dawsh 2y ago else:
117 d377c4c Dawsh 2y ago expire_strategy = "NEVER"
118 d377c4c Dawsh 2y ago connection.execute(
119 d377c4c Dawsh 2y ago users_table.update()
120 d377c4c Dawsh 2y ago .where(users_table.c.id == user[0])
121 d377c4c Dawsh 2y ago .values(expire_strategy=expire_strategy)
122 d377c4c Dawsh 2y ago )
123 d377c4c Dawsh 2y ago
124 d377c4c Dawsh 2y ago op.drop_column("users", "status")
125 d377c4c Dawsh 2y ago
126 d377c4c Dawsh 2y ago
127 d377c4c Dawsh 2y ago def downgrade() -> None:
128 d377c4c Dawsh 2y ago # ### commands auto generated by Alembic - please adjust! ###
129 d377c4c Dawsh 2y ago op.add_column(
130 d377c4c Dawsh 2y ago "users",
131 d377c4c Dawsh 2y ago sa.Column(
132 d377c4c Dawsh 2y ago "status",
133 d377c4c Dawsh 2y ago postgresql.ENUM(
134 d377c4c Dawsh 2y ago "active", "limited", "expired", "on_hold", name="userstatus"
135 d377c4c Dawsh 2y ago ),
136 d377c4c Dawsh 2y ago autoincrement=False,
137 d377c4c Dawsh 2y ago nullable=False,
138 d377c4c Dawsh 2y ago server_default="active",
139 d377c4c Dawsh 2y ago ),
140 d377c4c Dawsh 2y ago )
141 d377c4c Dawsh 2y ago op.alter_column(
142 deb4d55 Dawsh 2y ago "users",
143 deb4d55 Dawsh 2y ago "activation_deadline",
144 deb4d55 Dawsh 2y ago new_column_name="on_hold_timeout",
145 deb4d55 Dawsh 2y ago existing_type=sa.DateTime,
146 d377c4c Dawsh 2y ago )
147 d377c4c Dawsh 2y ago op.alter_column(
148 deb4d55 Dawsh 2y ago "users",
149 deb4d55 Dawsh 2y ago "usage_duration",
150 deb4d55 Dawsh 2y ago new_column_name="on_hold_expire_duration",
151 deb4d55 Dawsh 2y ago existing_type=sa.BigInteger,
152 deb4d55 Dawsh 2y ago )
153 deb4d55 Dawsh 2y ago op.alter_column(
154 deb4d55 Dawsh 2y ago "users",
155 deb4d55 Dawsh 2y ago "expire_date",
156 deb4d55 Dawsh 2y ago new_column_name="expire",
157 deb4d55 Dawsh 2y ago existing_type=sa.DateTime,
158 d377c4c Dawsh 2y ago )
159 d377c4c Dawsh 2y ago op.drop_column("users", "expire_strategy")
160 d377c4c Dawsh 2y ago op.drop_column("users", "activated")
161 d377c4c Dawsh 2y ago # ### end Alembic commands ###