From a952a1266be0a0d5ce233f93094543b30026814c Mon Sep 17 00:00:00 2001
From: Sarah Wooders <sarahwooders@gmail.com>
Date: Fri, 19 Sep 2025 11:23:17 -0700
Subject: [PATCH] feat: add various indexes (#4811)

---
 .../57bcea83af3f_add_various_indexes.py       | 43 +++++++++++++++++++
 letta/orm/block.py                            |  5 +++
 letta/orm/job.py                              |  5 ++-
 letta/orm/step.py                             |  3 +-
 4 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 alembic/versions/57bcea83af3f_add_various_indexes.py

diff --git a/alembic/versions/57bcea83af3f_add_various_indexes.py b/alembic/versions/57bcea83af3f_add_various_indexes.py
new file mode 100644
index 00000000..a2e71cd4
--- /dev/null
+++ b/alembic/versions/57bcea83af3f_add_various_indexes.py
@@ -0,0 +1,43 @@
+"""add various indexes
+
+Revision ID: 57bcea83af3f
+Revises: 5973fd8b8c60
+Create Date: 2025-09-19 10:58:19.658106
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "57bcea83af3f"
+down_revision: Union[str, None] = "5973fd8b8c60"
+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_block_hidden", "block", ["hidden"], unique=False)
+    op.create_index("ix_block_is_template", "block", ["is_template"], unique=False)
+    op.create_index("ix_block_org_project_template", "block", ["organization_id", "project_id", "is_template"], unique=False)
+    op.create_index("ix_block_organization_id", "block", ["organization_id"], unique=False)
+    op.create_index("ix_block_project_id", "block", ["project_id"], unique=False)
+    op.create_index("ix_jobs_user_id", "jobs", ["user_id"], unique=False)
+    op.create_index("ix_steps_job_id", "steps", ["job_id"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("ix_steps_job_id", table_name="steps")
+    op.drop_index("ix_jobs_user_id", table_name="jobs")
+    op.drop_index("ix_block_project_id", table_name="block")
+    op.drop_index("ix_block_organization_id", table_name="block")
+    op.drop_index("ix_block_org_project_template", table_name="block")
+    op.drop_index("ix_block_is_template", table_name="block")
+    op.drop_index("ix_block_hidden", table_name="block")
+    # ### end Alembic commands ###
diff --git a/letta/orm/block.py b/letta/orm/block.py
index bec65f4e..26023988 100644
--- a/letta/orm/block.py
+++ b/letta/orm/block.py
@@ -25,6 +25,11 @@ class Block(OrganizationMixin, SqlalchemyBase, ProjectMixin, TemplateEntityMixin
         UniqueConstraint("id", "label", name="unique_block_id_label"),
         Index("created_at_label_idx", "created_at", "label"),
         Index("ix_block_label", "label"),
+        Index("ix_block_organization_id", "organization_id"),
+        Index("ix_block_project_id", "project_id"),
+        Index("ix_block_is_template", "is_template"),
+        Index("ix_block_hidden", "hidden"),
+        Index("ix_block_org_project_template", "organization_id", "project_id", "is_template"),
     )
 
     template_name: Mapped[Optional[str]] = mapped_column(
diff --git a/letta/orm/job.py b/letta/orm/job.py
index 1b8981c8..048ea514 100644
--- a/letta/orm/job.py
+++ b/letta/orm/job.py
@@ -26,7 +26,10 @@ class Job(SqlalchemyBase, UserMixin):
 
     __tablename__ = "jobs"
     __pydantic_model__ = PydanticJob
-    __table_args__ = (Index("ix_jobs_created_at", "created_at", "id"),)
+    __table_args__ = (
+        Index("ix_jobs_created_at", "created_at", "id"),
+        Index("ix_jobs_user_id", "user_id"),
+    )
 
     status: Mapped[JobStatus] = mapped_column(String, default=JobStatus.created, doc="The current status of the job.")
     completed_at: Mapped[Optional[datetime]] = mapped_column(nullable=True, doc="The unix timestamp of when the job was completed.")
diff --git a/letta/orm/step.py b/letta/orm/step.py
index 49e90c42..444ab403 100644
--- a/letta/orm/step.py
+++ b/letta/orm/step.py
@@ -1,7 +1,7 @@
 import uuid
 from typing import TYPE_CHECKING, Dict, List, Optional
 
-from sqlalchemy import JSON, ForeignKey, String
+from sqlalchemy import JSON, ForeignKey, Index, String
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from letta.orm.mixins import ProjectMixin
@@ -22,6 +22,7 @@ class Step(SqlalchemyBase, ProjectMixin):
 
     __tablename__ = "steps"
     __pydantic_model__ = PydanticStep
+    __table_args__ = (Index("ix_steps_job_id", "job_id"),)
 
     id: Mapped[str] = mapped_column(String, primary_key=True, default=lambda: f"step-{uuid.uuid4()}")
     origin: Mapped[Optional[str]] = mapped_column(nullable=True, doc="The surface that this agent step was initiated from.")