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