fix: slow sql of ops tracing (#5749)

.github/workflows/db-migration-test.yml

@@ -38,6 +38,11 @@ jobs:
       - name: Install dependencies
         run: poetry install -C api
 
+      - name: Prepare middleware env
+        run: |
+          cd docker
+          cp middleware.env.example middleware.env
+
       - name: Set up Middlewares
         uses: hoverkraft-tech/compose-action@v2.0.0
         with:

api/core/ops/ops_trace_manager.py

@@ -352,10 +352,17 @@ class TraceTask:
         query = workflow_run_inputs.get("query") or workflow_run_inputs.get("sys.query") or ""
 
         # get workflow_app_log_id
-        workflow_app_log_data = db.session.query(WorkflowAppLog).filter_by(workflow_run_id=workflow_run.id).first()
+        workflow_app_log_data = db.session.query(WorkflowAppLog).filter_by(
+            tenant_id=tenant_id,
+            app_id=workflow_run.app_id,
+            workflow_run_id=workflow_run.id
+        ).first()
         workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None
         # get message_id
-        message_data = db.session.query(Message.id).filter_by(workflow_run_id=workflow_run_id).first()
+        message_data = db.session.query(Message.id).filter_by(
+            conversation_id=conversation_id,
+            workflow_run_id=workflow_run_id
+        ).first()
         message_id = str(message_data.id) if message_data else None
 
         metadata = {

api/migrations/versions/b2602e131636_add_workflow_run_id_index_for_message.py

@@ -0,0 +1,32 @@
+"""add workflow_run_id index for message
+
+Revision ID: b2602e131636
+Revises: 63f9175e515b
+Create Date: 2024-06-29 12:16:51.646346
+
+"""
+from alembic import op
+
+import models as models
+
+# revision identifiers, used by Alembic.
+revision = 'b2602e131636'
+down_revision = '63f9175e515b'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('messages', schema=None) as batch_op:
+        batch_op.create_index('message_workflow_run_id_idx', ['conversation_id', 'workflow_run_id'], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('messages', schema=None) as batch_op:
+        batch_op.drop_index('message_workflow_run_id_idx')
+
+    # ### end Alembic commands ###

api/models/model.py

@@ -626,6 +626,7 @@ class Message(db.Model):
         db.Index('message_conversation_id_idx', 'conversation_id'),
         db.Index('message_end_user_idx', 'app_id', 'from_source', 'from_end_user_id'),
         db.Index('message_account_idx', 'app_id', 'from_source', 'from_account_id'),
+        db.Index('message_workflow_run_id_idx', 'conversation_id', 'workflow_run_id')
     )
 
     id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))