add tenant id index for document and document_segment table (#2873)

Co-authored-by: jyong <jyong@dify.ai>

api/migrations/versions/a8f9b3c45e4a_add_tenant_id_db_index.py

@@ -0,0 +1,36 @@
+"""add_tenant_id_db_index
+
+Revision ID: a8f9b3c45e4a
+Revises: 16830a790f0f
+Create Date: 2024-03-18 05:07:35.588473
+
+"""
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'a8f9b3c45e4a'
+down_revision = '16830a790f0f'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('document_segments', schema=None) as batch_op:
+        batch_op.create_index('document_segment_tenant_idx', ['tenant_id'], unique=False)
+
+    with op.batch_alter_table('documents', schema=None) as batch_op:
+        batch_op.create_index('document_tenant_idx', ['tenant_id'], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('documents', schema=None) as batch_op:
+        batch_op.drop_index('document_tenant_idx')
+
+    with op.batch_alter_table('document_segments', schema=None) as batch_op:
+        batch_op.drop_index('document_segment_tenant_idx')
+
+    # ### end Alembic commands ###

api/models/dataset.py

@@ -176,6 +176,7 @@ class Document(db.Model):
         db.PrimaryKeyConstraint('id', name='document_pkey'),
         db.Index('document_dataset_id_idx', 'dataset_id'),
         db.Index('document_is_paused_idx', 'is_paused'),
+        db.Index('document_tenant_idx', 'tenant_id'),
     )
 
     # initial fields
@@ -334,6 +335,7 @@ class DocumentSegment(db.Model):
         db.Index('document_segment_tenant_dataset_idx', 'dataset_id', 'tenant_id'),
         db.Index('document_segment_tenant_document_idx', 'document_id', 'tenant_id'),
         db.Index('document_segment_dataset_node_idx', 'dataset_id', 'index_node_id'),
+        db.Index('document_segment_tenant_idx', 'tenant_id'),
     )
 
     # initial fields