Skip to content

Commit 36ed9b7

Browse files
Mary Hippmaryhipp
authored andcommitted
restore list_queue_items method
1 parent 3cec06f commit 36ed9b7

File tree

2 files changed

+69
-0
lines changed

2 files changed

+69
-0
lines changed

invokeai/app/services/session_queue/session_queue_base.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@
2020
SessionQueueCountsByDestination,
2121
SessionQueueItem,
2222
SessionQueueStatus,
23+
QUEUE_ITEM_STATUS,
2324
)
25+
from invokeai.app.services.shared.pagination import CursorPaginatedResults
2426
from invokeai.app.services.shared.graph import GraphExecutionState
2527
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
2628

@@ -135,6 +137,19 @@ def delete_all_except_current(self, queue_id: str) -> DeleteAllExceptCurrentResu
135137
"""Deletes all queue items except in-progress items"""
136138
pass
137139

140+
@abstractmethod
141+
def list_queue_items(
142+
self,
143+
queue_id: str,
144+
limit: int,
145+
priority: int,
146+
cursor: Optional[int] = None,
147+
status: Optional[QUEUE_ITEM_STATUS] = None,
148+
destination: Optional[str] = None,
149+
) -> CursorPaginatedResults[SessionQueueItem]:
150+
"""Gets a page of session queue items. Do not remove."""
151+
pass
152+
138153
@abstractmethod
139154
def list_all_queue_items(
140155
self,

invokeai/app/services/session_queue/session_queue_sqlite.py

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import sqlite3
44
from typing import Optional, Union, cast
55

6+
from invokeai.app.services.shared.pagination import CursorPaginatedResults
67
from pydantic_core import to_jsonable_python
78

89
from invokeai.app.services.invoker import Invoker
@@ -588,6 +589,59 @@ def set_queue_item_session(self, item_id: int, session: GraphExecutionState) ->
588589
)
589590
return self.get_queue_item(item_id)
590591

592+
def list_queue_items(
593+
self,
594+
queue_id: str,
595+
limit: int,
596+
priority: int,
597+
cursor: Optional[int] = None,
598+
status: Optional[QUEUE_ITEM_STATUS] = None,
599+
destination: Optional[str] = None,
600+
) -> CursorPaginatedResults[SessionQueueItem]:
601+
with self._db.transaction() as cursor_:
602+
item_id = cursor
603+
query = """--sql
604+
SELECT *
605+
FROM session_queue
606+
WHERE queue_id = ?
607+
"""
608+
params: list[Union[str, int]] = [queue_id]
609+
610+
if status is not None:
611+
query += """--sql
612+
AND status = ?
613+
"""
614+
params.append(status)
615+
616+
if destination is not None:
617+
query += """---sql
618+
AND destination = ?
619+
"""
620+
params.append(destination)
621+
622+
if item_id is not None:
623+
query += """--sql
624+
AND (priority < ?) OR (priority = ? AND item_id > ?)
625+
"""
626+
params.extend([priority, priority, item_id])
627+
628+
query += """--sql
629+
ORDER BY
630+
priority DESC,
631+
item_id ASC
632+
LIMIT ?
633+
"""
634+
params.append(limit + 1)
635+
cursor_.execute(query, params)
636+
results = cast(list[sqlite3.Row], cursor_.fetchall())
637+
items = [SessionQueueItem.queue_item_from_dict(dict(result)) for result in results]
638+
has_more = False
639+
if len(items) > limit:
640+
# remove the extra item
641+
items.pop()
642+
has_more = True
643+
return CursorPaginatedResults(items=items, limit=limit, has_more=has_more)
644+
591645
def list_all_queue_items(
592646
self,
593647
queue_id: str,

0 commit comments

Comments
 (0)