Improve the demo feature of the project

services/demo_session/migrations/versions/002_add_cloning_status_tracking.py

@@ -0,0 +1,81 @@
+"""Add cloning status tracking
+
+Revision ID: 002
+Revises: 001
+Create Date: 2025-01-10
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers
+revision = '002'
+down_revision = 'a1b2c3d4e5f6'  # References the actual initial schema revision
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    """Add new status values and cloning tracking fields"""
+
+    # Add new columns for cloning progress
+    op.add_column('demo_sessions', sa.Column('cloning_started_at', sa.DateTime(timezone=True), nullable=True))
+    op.add_column('demo_sessions', sa.Column('cloning_completed_at', sa.DateTime(timezone=True), nullable=True))
+    op.add_column('demo_sessions', sa.Column('total_records_cloned', sa.Integer(), server_default='0', nullable=False))
+    op.add_column('demo_sessions', sa.Column('cloning_progress', postgresql.JSONB(astext_type=sa.Text()), server_default='{}', nullable=False))
+
+    # Update the status enum to include new values
+    # PostgreSQL doesn't support IF NOT EXISTS for enum values in older versions
+    # We need to check if values exist before adding them
+    from sqlalchemy import text
+
+    conn = op.get_bind()
+
+    # Check and add each enum value if it doesn't exist
+    enum_values_to_add = ['pending', 'ready', 'failed', 'partial']
+
+    for value in enum_values_to_add:
+        # Check if the enum value already exists
+        result = conn.execute(text("""
+            SELECT EXISTS (
+                SELECT 1 FROM pg_enum
+                WHERE enumlabel = :value
+                AND enumtypid = (
+                    SELECT oid FROM pg_type WHERE typname = 'demosessionstatus'
+                )
+            );
+        """), {"value": value})
+
+        exists = result.scalar()
+
+        if not exists:
+            # Add the enum value
+            # Note: ALTER TYPE ADD VALUE cannot run inside a transaction block in PostgreSQL
+            # but Alembic handles this for us
+            conn.execute(text(f"ALTER TYPE demosessionstatus ADD VALUE '{value}'"))
+
+    # Update existing sessions: active → ready
+    op.execute("""
+        UPDATE demo_sessions
+        SET status = 'ready'
+        WHERE status = 'active' AND data_cloned = true;
+    """)
+
+
+def downgrade():
+    """Remove cloning status tracking"""
+
+    # Remove new columns
+    op.drop_column('demo_sessions', 'cloning_progress')
+    op.drop_column('demo_sessions', 'total_records_cloned')
+    op.drop_column('demo_sessions', 'cloning_completed_at')
+    op.drop_column('demo_sessions', 'cloning_started_at')
+
+    # Note: Cannot easily remove enum values in PostgreSQL
+    # Migration down would require recreating the enum type
+    op.execute("""
+        UPDATE demo_sessions
+        SET status = 'active'
+        WHERE status IN ('ready', 'pending', 'failed', 'partial');
+    """)