Fix migration

hasgeek · Nov 20, 2023 · 164ac21 · 164ac21
1 parent e905764
commit 164ac21
Showing 1 changed file with 15 additions and 4 deletions.
diff --git a/migrations/versions/f0ed25eed4bc_replace_rsvp_flag.py b/migrations/versions/f0ed25eed4bc_replace_rsvp_flag.py
@@ -27,6 +27,7 @@
 class RsvpState:
     NONE = 1
     ALL = 2
+    MEMBERS = 3
 
 
 def upgrade(engine_name: str = '') -> None:
@@ -43,7 +44,17 @@ def downgrade(engine_name: str = '') -> None:
 
 def upgrade_() -> None:
     """Upgrade default database."""
-    op.add_column('project', sa.Column('rsvp_state', sa.SmallInteger(), nullable=True))
+    op.add_column(
+        'project',
+        sa.Column(
+            'rsvp_state',
+            sa.SmallInteger(),
+            sa.CheckConstraint(
+                'rsvp_state IN (1, 2, 3)', name='project_rsvp_state_check'
+            ),
+            nullable=True,
+        ),
+    )
     op.execute(
         project.update().values(
             rsvp_state=sa.case(
@@ -65,9 +76,9 @@ def downgrade_() -> None:
     )
     op.execute(
         project.update().values(
-            rsvp_state=sa.case(
-                (project.c.rsvp_state.is_(RsvpState.NONE), False),
-                (project.c.rsvp_state.is_(RsvpState.ALL), True),
+            allow_rsvp=sa.case(
+                (project.c.rsvp_state == RsvpState.NONE, False),
+                (project.c.rsvp_state == RsvpState.ALL, True),
                 else_=False,
             )
         )