update migrate script (#191)

opendatalab · Feb 21, 2025 · 3591bcf · 3591bcf
1 parent c436b00
commit 3591bcf
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 41 deletions.
diff --git a/labelu/alembic_labelu/versions/2eb983c9a254_add_collaborators_and_updaters.py b/labelu/alembic_labelu/versions/2eb983c9a254_add_collaborators_and_updaters.py
@@ -20,47 +20,59 @@
 
 def upgrade() -> None:
     # Create task_collaborator table
-    op.create_table(
-        'task_collaborator',
-        sa.Column('task_id', sa.Integer(), nullable=False),
-        sa.Column('user_id', sa.Integer(), nullable=False),
-        sa.Column(
-            'created_at',
-            sa.DateTime(timezone=True),
-            server_default=sa.text('CURRENT_TIMESTAMP'),
-            nullable=False
-        ),
-        sa.ForeignKeyConstraint(
-            ['task_id'],
-            ['task.id'],
-            ondelete='CASCADE'
-        ),
-        sa.ForeignKeyConstraint(
-            ['user_id'],
-            ['user.id'],
-            ondelete='CASCADE'
-        ),
-        sa.PrimaryKeyConstraint('task_id', 'user_id')
-    )
-
-    # Performances index
-    op.create_index(
-        'ix_task_collaborator_task_id',
-        'task_collaborator',
-        ['task_id']
-    )
-    op.create_index(
-        'ix_task_collaborator_user_id',
-        'task_collaborator',
-        ['user_id']
-    )
-    op.create_index(
-        'ix_task_created_by_deleted_at',
-        'task',
-        ['created_by', 'deleted_at']
-    )
+    # if the table is not exists then create it
+    is_table_exist = op.execute(
+        "SHOW TABLES LIKE 'task_collaborator';"
+    ).fetchone()
+    if not is_table_exist:
+        op.create_table(
+            'task_collaborator',
+            sa.Column('task_id', sa.Integer(), nullable=False),
+            sa.Column('user_id', sa.Integer(), nullable=False),
+            sa.Column(
+                'created_at',
+                sa.DateTime(timezone=True),
+                server_default=sa.text('CURRENT_TIMESTAMP'),
+                nullable=False
+            ),
+            sa.ForeignKeyConstraint(
+                ['task_id'],
+                ['task.id'],
+                ondelete='CASCADE'
+            ),
+            sa.ForeignKeyConstraint(
+                ['user_id'],
+                ['user.id'],
+                ondelete='CASCADE'
+            ),
+            sa.PrimaryKeyConstraint('task_id', 'user_id')
+        )
+
+        # Performances index
+        op.create_index(
+            'ix_task_collaborator_task_id',
+            'task_collaborator',
+            ['task_id']
+        )
+        op.create_index(
+            'ix_task_collaborator_user_id',
+            'task_collaborator',
+            ['user_id']
+        )
+        op.create_index(
+            'ix_task_created_by_deleted_at',
+            'task',
+            ['created_by', 'deleted_at']
+        )
 
     # Task sample: updater -> updaters; create a new table task_sample_updater
+    is_task_sample_updater_table_exist = op.execute(
+        "SHOW TABLES LIKE 'task_sample_updater';"
+    ).fetchone()
+
+    if is_task_sample_updater_table_exist:
+        return
+
     op.create_table(
         'task_sample_updater',
         sa.Column('sample_id', sa.Integer(), nullable=False),

diff --git a/labelu/internal/domain/models/task_collaborator.py b/labelu/internal/domain/models/task_collaborator.py
@@ -1,4 +1,5 @@
 from datetime import datetime
+from sqlalchemy.schema import Index
 from sqlalchemy import Column, DateTime, ForeignKey, Integer
 from labelu.internal.common.db import Base
 
@@ -21,4 +22,8 @@ class TaskCollaborator(Base):
         DateTime(timezone=True),
         default=datetime.now,
         comment="Time a task collaborator was created"
-    )
+    )
+
+    Index("ix_task_collaborator_task_id", task_id)
+    Index("ix_task_collaborator_user_id", user_id)
+    Index("ix_task_created_by_deleted_at", task_id, user_id)
diff --git a/labelu/internal/domain/models/task_sample_updater.py b/labelu/internal/domain/models/task_sample_updater.py
@@ -1,4 +1,5 @@
 from datetime import datetime
+from sqlalchemy.schema import Index
 from sqlalchemy import Column, DateTime, ForeignKey, Integer
 from labelu.internal.common.db import Base
 
@@ -23,4 +24,7 @@ class TaskSampleUpdater(Base):
         onupdate=datetime.now,
         nullable=False,
         comment="created time"
-    )
+    )
+
+    Index("ix_task_sample_updater_sample_id", sample_id)
+    Index("ix_task_sample_updater_user_id", user_id)