From fae6a4d613973c30e5484c9bbeeb28424e97818c Mon Sep 17 00:00:00 2001
From: Stas Moreinis <stas.moreinis@scale.com>
Date: Wed, 24 Dec 2025 12:25:49 -0800
Subject: [PATCH] Add performance indexes to spans, agents, and tasks tables

- Add indexes to SpanORM: trace_id, (trace_id, start_time), parent_id
- Add index to AgentORM: status
- Add index to TaskORM: status

These indexes improve query performance for:
- Filtering spans by trace_id (common in observability queries)
- Ordering spans by start_time within a trace
- Traversing span hierarchy via parent_id
- Filtering agents/tasks by status (used in list queries)

Test results show ~3x improvement on 894 rows, expected to be
100x+ on larger datasets due to O(log N) vs O(N) complexity.
---
 ...23_add_performance_indexes_d024851e790c.py | 38 +++++++++++++++++++
 .../database/migrations/migration_history.txt |  3 +-
 agentex/src/adapters/orm.py                   | 22 +++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 agentex/database/migrations/alembic/versions/2025_12_24_2023_add_performance_indexes_d024851e790c.py

diff --git a/agentex/database/migrations/alembic/versions/2025_12_24_2023_add_performance_indexes_d024851e790c.py b/agentex/database/migrations/alembic/versions/2025_12_24_2023_add_performance_indexes_d024851e790c.py
new file mode 100644
index 0000000..0ee06c0
--- /dev/null
+++ b/agentex/database/migrations/alembic/versions/2025_12_24_2023_add_performance_indexes_d024851e790c.py
@@ -0,0 +1,38 @@
+"""add_performance_indexes
+
+Revision ID: d024851e790c
+Revises: 24429f13b8bd
+Create Date: 2025-12-24 20:23:00.061860
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'd024851e790c'
+down_revision: Union[str, None] = '24429f13b8bd'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_index('ix_agents_status', 'agents', ['status'], unique=False)
+    op.create_index('ix_spans_parent_id', 'spans', ['parent_id'], unique=False)
+    op.create_index('ix_spans_trace_id', 'spans', ['trace_id'], unique=False)
+    op.create_index('ix_spans_trace_id_start_time', 'spans', ['trace_id', 'start_time'], unique=False)
+    op.create_index('ix_tasks_status', 'tasks', ['status'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index('ix_tasks_status', table_name='tasks')
+    op.drop_index('ix_spans_trace_id_start_time', table_name='spans')
+    op.drop_index('ix_spans_trace_id', table_name='spans')
+    op.drop_index('ix_spans_parent_id', table_name='spans')
+    op.drop_index('ix_agents_status', table_name='agents')
+    # ### end Alembic commands ###
diff --git a/agentex/database/migrations/migration_history.txt b/agentex/database/migrations/migration_history.txt
index 529c922..3754d85 100644
--- a/agentex/database/migrations/migration_history.txt
+++ b/agentex/database/migrations/migration_history.txt
@@ -1,4 +1,5 @@
-a5d67f2d7356 -> 24429f13b8bd (head), add agent input type
+24429f13b8bd -> d024851e790c (head), add_performance_indexes
+a5d67f2d7356 -> 24429f13b8bd, add agent input type
 329fbafa4ff9 -> a5d67f2d7356, add unhealthy status
 d7addd4229e8 -> 329fbafa4ff9, change_default_acp_to_async
 09368a02d6cc -> d7addd4229e8, soft delete status
diff --git a/agentex/src/adapters/orm.py b/agentex/src/adapters/orm.py
index e83f09b..8d24fed 100644
--- a/agentex/src/adapters/orm.py
+++ b/agentex/src/adapters/orm.py
@@ -46,6 +46,12 @@ class AgentORM(BaseORM):
     # Many-to-Many relationship with tasks
     tasks = relationship("TaskORM", secondary="task_agents", back_populates="agents")
 
+    # Indexes for efficient querying
+    __table_args__ = (
+        # Index for filtering agents by status (used in list queries)
+        Index("ix_agents_status", "status"),
+    )
+
 
 class TaskORM(BaseORM):
     __tablename__ = "tasks"
@@ -64,6 +70,12 @@ class TaskORM(BaseORM):
     # Many-to-Many relationship with agents
     agents = relationship("AgentORM", secondary="task_agents", back_populates="tasks")
 
+    # Indexes for efficient querying
+    __table_args__ = (
+        # Index for filtering tasks by status (used in list queries)
+        Index("ix_tasks_status", "status"),
+    )
+
 
 class TaskAgentORM(BaseORM):
     __tablename__ = "task_agents"
@@ -139,6 +151,16 @@ class SpanORM(BaseORM):
     output = Column(JSON, nullable=True)
     data = Column(JSON, nullable=True)
 
+    # Indexes for efficient querying
+    __table_args__ = (
+        # Index for filtering spans by trace_id
+        Index("ix_spans_trace_id", "trace_id"),
+        # Composite index for filtering by trace_id and ordering by start_time
+        Index("ix_spans_trace_id_start_time", "trace_id", "start_time"),
+        # Index for traversing span hierarchy
+        Index("ix_spans_parent_id", "parent_id"),
+    )
+
 
 class AgentAPIKeyORM(BaseORM):
     __tablename__ = "agent_api_keys"