|
3 | 3 | import sqlite3
|
4 | 4 | from typing import Optional, Union, cast
|
5 | 5 |
|
| 6 | +from invokeai.app.services.shared.pagination import CursorPaginatedResults |
6 | 7 | from pydantic_core import to_jsonable_python
|
7 | 8 |
|
8 | 9 | from invokeai.app.services.invoker import Invoker
|
@@ -588,6 +589,59 @@ def set_queue_item_session(self, item_id: int, session: GraphExecutionState) ->
|
588 | 589 | )
|
589 | 590 | return self.get_queue_item(item_id)
|
590 | 591 |
|
| 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 | + |
591 | 645 | def list_all_queue_items(
|
592 | 646 | self,
|
593 | 647 | queue_id: str,
|
|
0 commit comments