From 1e045a0187d2e637c6156a78c12778b842b89ab1 Mon Sep 17 00:00:00 2001 From: takatost Date: Sat, 29 Jun 2024 20:28:30 +0800 Subject: [PATCH] fix: slow sql of ops tracing (#5749) --- .github/workflows/db-migration-test.yml | 5 +++ api/core/ops/ops_trace_manager.py | 11 +++++-- ...6_add_workflow_run_id_index_for_message.py | 32 +++++++++++++++++++ api/models/model.py | 1 + 4 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 api/migrations/versions/b2602e131636_add_workflow_run_id_index_for_message.py diff --git a/.github/workflows/db-migration-test.yml b/.github/workflows/db-migration-test.yml index 64e8eb291c..67d1558dbc 100644 --- a/.github/workflows/db-migration-test.yml +++ b/.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: diff --git a/api/core/ops/ops_trace_manager.py b/api/core/ops/ops_trace_manager.py index fbdc12331e..ff15aa999b 100644 --- a/api/core/ops/ops_trace_manager.py +++ b/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 = { diff --git a/api/migrations/versions/b2602e131636_add_workflow_run_id_index_for_message.py b/api/migrations/versions/b2602e131636_add_workflow_run_id_index_for_message.py new file mode 100644 index 0000000000..c9a6a5a5a7 --- /dev/null +++ b/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 ### diff --git a/api/models/model.py b/api/models/model.py index 07d7f6d891..f59e8ebb7c 100644 --- a/api/models/model.py +++ b/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()'))