Add indexes to improve query performance on various tables (#4955)

Author: A-Artemis

application/backend/app/alembic/versions/2786b50eb5a4_schema.py

@@ -31,6 +31,7 @@ def upgrade() -> None:
         sa.Column("updated_at", sa.DateTime(), server_default=sa.text("(CURRENT_TIMESTAMP)"), nullable=False),
         sa.PrimaryKeyConstraint("id"),
     )
+    op.create_index("idx_projects_name", "projects", ["name"], unique=False)
     op.create_table(
         "sinks",
         sa.Column("id", sa.Text(), nullable=False),
@@ -65,6 +66,7 @@ def upgrade() -> None:
         sa.ForeignKeyConstraint(["project_id"], ["projects.id"], ondelete="CASCADE"),
         sa.PrimaryKeyConstraint("id"),
     )
+    op.create_index("idx_dataset_revisions_project", "dataset_revisions", ["project_id"], unique=False)
     op.create_table(
         "labels",
         sa.Column("project_id", sa.Text(), nullable=False),
@@ -94,17 +96,15 @@ def upgrade() -> None:
         sa.Column("id", sa.Text(), nullable=False),
         sa.Column("created_at", sa.DateTime(), server_default=sa.text("(CURRENT_TIMESTAMP)"), nullable=False),
         sa.Column("updated_at", sa.DateTime(), server_default=sa.text("(CURRENT_TIMESTAMP)"), nullable=False),
-        sa.ForeignKeyConstraint(
-            ["parent_revision"],
-            ["model_revisions.id"],
-        ),
+        sa.ForeignKeyConstraint(["parent_revision"], ["model_revisions.id"]),
         sa.ForeignKeyConstraint(["project_id"], ["projects.id"], ondelete="CASCADE"),
-        sa.ForeignKeyConstraint(
-            ["training_dataset_id"],
-            ["dataset_revisions.id"],
-        ),
+        sa.ForeignKeyConstraint(["training_dataset_id"], ["dataset_revisions.id"]),
         sa.PrimaryKeyConstraint("id"),
     )
+    op.create_index(
+        "idx_model_revisions_project_status", "model_revisions", ["project_id", "training_status"], unique=False
+    )
+    op.create_index("idx_model_revisions_architecture", "model_revisions", ["project_id", "architecture"], unique=False)
     op.create_table(
         "dataset_items",
         sa.Column("project_id", sa.Text(), nullable=False),
@@ -127,6 +127,7 @@ def upgrade() -> None:
         sa.ForeignKeyConstraint(["source_id"], ["sources.id"], ondelete="SET NULL"),
         sa.PrimaryKeyConstraint("id"),
     )
+    op.create_index("idx_dataset_items_user_reviewed", "dataset_items", ["project_id", "user_reviewed"], unique=False)
     op.create_index("idx_dataset_items_project_created_at", "dataset_items", ["project_id", "created_at"], unique=False)
     op.create_table(
         "pipelines",
@@ -144,6 +145,7 @@ def upgrade() -> None:
         sa.ForeignKeyConstraint(["source_id"], ["sources.id"], ondelete="RESTRICT"),
         sa.PrimaryKeyConstraint("project_id"),
     )
+    op.create_index("idx_pipelines_is_running", "pipelines", ["is_running"], unique=False)
     op.create_table(
         "dataset_items_labels",
         sa.Column("dataset_item_id", sa.Text(), nullable=False),
@@ -161,13 +163,19 @@ def downgrade() -> None:
     """Downgrade schema."""
     # ### commands auto generated by Alembic - please adjust! ###
     op.drop_table("dataset_items_labels")
+    op.drop_index("idx_pipelines_is_running", table_name="pipelines")
     op.drop_table("pipelines")
     op.drop_index("idx_dataset_items_project_created_at", table_name="dataset_items")
+    op.drop_index("idx_dataset_items_user_reviewed", table_name="dataset_items")
     op.drop_table("dataset_items")
+    op.drop_index("idx_model_revisions_architecture", table_name="model_revisions")
+    op.drop_index("idx_model_revisions_project_status", table_name="model_revisions")
     op.drop_table("model_revisions")
     op.drop_table("labels")
+    op.drop_index("idx_dataset_revisions_project", table_name="dataset_revisions")
     op.drop_table("dataset_revisions")
     op.drop_table("sources")
     op.drop_table("sinks")
+    op.drop_index("idx_projects_name", table_name="projects")
     op.drop_table("projects")
     # ### end Alembic commands ###

application/backend/app/db/schema.py

@@ -29,6 +29,7 @@ class SourceDB(BaseID):
 
 class ProjectDB(BaseID):
     __tablename__ = "projects"
+    __table_args__ = (Index("idx_projects_name", "name"),)
 
     name: Mapped[str] = mapped_column(String(255), nullable=False)
     task_type: Mapped[str] = mapped_column(String(50), nullable=False)
@@ -39,6 +40,7 @@ class ProjectDB(BaseID):
 
 class PipelineDB(Base):
     __tablename__ = "pipelines"
+    __table_args__ = (Index("idx_pipelines_is_running", "is_running"),)
 
     project_id: Mapped[str] = mapped_column(Text, ForeignKey("projects.id", ondelete="CASCADE"), primary_key=True)
     source_id: Mapped[str | None] = mapped_column(Text, ForeignKey("sources.id", ondelete="RESTRICT"))
@@ -65,6 +67,10 @@ class SinkDB(BaseID):
 
 class ModelRevisionDB(BaseID):
     __tablename__ = "model_revisions"
+    __table_args__ = (
+        Index("idx_model_revisions_project_status", "project_id", "training_status"),
+        Index("idx_model_revisions_architecture", "project_id", "architecture"),
+    )
 
     project_id: Mapped[str] = mapped_column(Text, ForeignKey("projects.id", ondelete="CASCADE"), nullable=False)
     architecture: Mapped[str] = mapped_column(String(100), nullable=False)
@@ -82,14 +88,15 @@ class ModelRevisionDB(BaseID):
 
 class DatasetRevisionDB(BaseID):
     __tablename__ = "dataset_revisions"
+    __table_args__ = (Index("idx_dataset_revisions_project", "project_id"),)
 
     project_id: Mapped[str] = mapped_column(Text, ForeignKey("projects.id", ondelete="CASCADE"), nullable=False)
     files_deleted: Mapped[bool] = mapped_column(Boolean, default=False)
 
 
 class DatasetItemDB(BaseID):
     __tablename__ = "dataset_items"
-    __table_args__ = (Index("idx_dataset_items_project_created_at", "project_id", "created_at"),)
+    __table_args__ = (Index("idx_dataset_items_user_reviewed", "project_id", "user_reviewed"),)
 
     project_id: Mapped[str] = mapped_column(Text, ForeignKey("projects.id", ondelete="CASCADE"), nullable=False)
     name: Mapped[str] = mapped_column(String(255), nullable=False)