refactor and optimize get_project_states_in_between

This commit is contained in:
LeonOstrez
2025-07-08 21:41:16 +02:00
parent 2f51ffc88c
commit 91bf445efa
2 changed files with 24 additions and 9 deletions

View File

@@ -708,7 +708,9 @@ class ProjectState(Base):
return epics_and_tasks
@staticmethod
async def get_project_states_in_between(session: "AsyncSession", branch_id: UUID, start_id: UUID, end_id: UUID):
async def get_project_states_in_between(
session: "AsyncSession", branch_id: UUID, start_id: UUID, end_id: UUID, limit: Optional[int] = 100
):
query = select(ProjectState).where(
and_(
ProjectState.branch_id == branch_id,
@@ -731,15 +733,26 @@ class ProjectState(Base):
log.error(f"Could not find states with IDs {start_id} and {end_id} in branch {branch_id}")
return []
query = select(ProjectState).where(
and_(
ProjectState.branch_id == branch_id,
ProjectState.step_index >= start_state.step_index,
ProjectState.step_index <= end_state.step_index,
query = (
select(ProjectState)
.where(
and_(
ProjectState.branch_id == branch_id,
ProjectState.step_index >= start_state.step_index,
ProjectState.step_index <= end_state.step_index,
)
)
.order_by(ProjectState.step_index.desc())
)
if limit:
query = query.limit(limit)
result = await session.execute(query)
return result.scalars().all()
states = result.scalars().all()
# Since we always order by step_index desc, we need to reverse to get chronological order
return list(reversed(states))
@staticmethod
async def get_task_conversation_project_states(

View File

@@ -148,13 +148,15 @@ class StateManager:
self.current_session, self.current_state.branch_id, task_id, first_last_only, limit
)
async def get_project_states_in_between(self, start_state_id: UUID, end_state_id: UUID) -> list[ProjectState]:
async def get_project_states_in_between(
self, start_state_id: UUID, end_state_id: UUID, limit: Optional[int] = 100
) -> list[ProjectState]:
"""
Get all project states in between two states.
This retrieves all project states that are associated with a specific branch
"""
return await ProjectState.get_project_states_in_between(
self.current_session, self.current_state.branch_id, start_state_id, end_state_id
self.current_session, self.current_state.branch_id, start_state_id, end_state_id, limit
)
async def get_fe_states(self) -> Optional[ProjectState]: