|
19 | 19 | from invokeai.app.services.shared.pagination import OffsetPaginatedResults |
20 | 20 | from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection |
21 | 21 | from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase |
| 22 | +from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO |
22 | 23 |
|
23 | 24 |
|
24 | 25 | class SqliteImageRecordStorage(ImageRecordStorageBase): |
@@ -503,3 +504,141 @@ def get_image_names( |
503 | 504 | image_names = [row[0] for row in result] |
504 | 505 |
|
505 | 506 | return ImageNamesResult(image_names=image_names, starred_count=starred_count, total_count=len(image_names)) |
| 507 | + |
| 508 | + def get_image_dates( |
| 509 | + self, |
| 510 | + user_id: Optional[str] = None, |
| 511 | + is_admin: bool = False, |
| 512 | + ) -> list[VirtualSubBoardDTO]: |
| 513 | + with self._db.transaction() as cursor: |
| 514 | + query_conditions = "" |
| 515 | + query_params: list[Union[int, str, bool]] = [] |
| 516 | + |
| 517 | + # Only non-intermediate images |
| 518 | + query_conditions += """--sql |
| 519 | + AND images.is_intermediate = 0 |
| 520 | + """ |
| 521 | + |
| 522 | + # User isolation for non-admin users |
| 523 | + if user_id is not None and not is_admin: |
| 524 | + query_conditions += """--sql |
| 525 | + AND images.user_id = ? |
| 526 | + """ |
| 527 | + query_params.append(user_id) |
| 528 | + |
| 529 | + query = f"""--sql |
| 530 | + SELECT |
| 531 | + DATE(images.created_at) as date, |
| 532 | + SUM(CASE WHEN images.image_category = 'general' THEN 1 ELSE 0 END) as image_count, |
| 533 | + SUM(CASE WHEN images.image_category != 'general' THEN 1 ELSE 0 END) as asset_count, |
| 534 | + ( |
| 535 | + SELECT i2.image_name FROM images i2 |
| 536 | + WHERE DATE(i2.created_at) = DATE(images.created_at) |
| 537 | + AND i2.is_intermediate = 0 |
| 538 | + ORDER BY i2.created_at DESC LIMIT 1 |
| 539 | + ) as cover_image_name |
| 540 | + FROM images |
| 541 | + WHERE 1=1 |
| 542 | + {query_conditions} |
| 543 | + GROUP BY DATE(images.created_at) |
| 544 | + ORDER BY date DESC; |
| 545 | + """ |
| 546 | + |
| 547 | + cursor.execute(query, query_params) |
| 548 | + result = cast(list[sqlite3.Row], cursor.fetchall()) |
| 549 | + |
| 550 | + return [ |
| 551 | + VirtualSubBoardDTO( |
| 552 | + virtual_board_id=f"by_date:{dict(row)['date']}", |
| 553 | + board_name=dict(row)["date"], |
| 554 | + date=dict(row)["date"], |
| 555 | + image_count=dict(row)["image_count"], |
| 556 | + asset_count=dict(row)["asset_count"], |
| 557 | + cover_image_name=dict(row)["cover_image_name"], |
| 558 | + ) |
| 559 | + for row in result |
| 560 | + ] |
| 561 | + |
| 562 | + def get_image_names_by_date( |
| 563 | + self, |
| 564 | + date: str, |
| 565 | + starred_first: bool = True, |
| 566 | + order_dir: SQLiteDirection = SQLiteDirection.Descending, |
| 567 | + categories: Optional[list[ImageCategory]] = None, |
| 568 | + search_term: Optional[str] = None, |
| 569 | + user_id: Optional[str] = None, |
| 570 | + is_admin: bool = False, |
| 571 | + ) -> ImageNamesResult: |
| 572 | + with self._db.transaction() as cursor: |
| 573 | + query_conditions = "" |
| 574 | + query_params: list[Union[int, str, bool]] = [] |
| 575 | + |
| 576 | + # Filter by date |
| 577 | + query_conditions += """--sql |
| 578 | + AND DATE(images.created_at) = ? |
| 579 | + """ |
| 580 | + query_params.append(date) |
| 581 | + |
| 582 | + # Only non-intermediate images |
| 583 | + query_conditions += """--sql |
| 584 | + AND images.is_intermediate = 0 |
| 585 | + """ |
| 586 | + |
| 587 | + if categories is not None: |
| 588 | + category_strings = [c.value for c in set(categories)] |
| 589 | + placeholders = ",".join("?" * len(category_strings)) |
| 590 | + query_conditions += f"""--sql |
| 591 | + AND images.image_category IN ( {placeholders} ) |
| 592 | + """ |
| 593 | + for c in category_strings: |
| 594 | + query_params.append(c) |
| 595 | + |
| 596 | + # User isolation for non-admin users |
| 597 | + if user_id is not None and not is_admin: |
| 598 | + query_conditions += """--sql |
| 599 | + AND images.user_id = ? |
| 600 | + """ |
| 601 | + query_params.append(user_id) |
| 602 | + |
| 603 | + if search_term: |
| 604 | + query_conditions += """--sql |
| 605 | + AND ( |
| 606 | + images.metadata LIKE ? |
| 607 | + OR images.created_at LIKE ? |
| 608 | + ) |
| 609 | + """ |
| 610 | + query_params.append(f"%{search_term.lower()}%") |
| 611 | + query_params.append(f"%{search_term.lower()}%") |
| 612 | + |
| 613 | + # Get starred count if starred_first is enabled |
| 614 | + starred_count = 0 |
| 615 | + if starred_first: |
| 616 | + starred_count_query = f"""--sql |
| 617 | + SELECT COUNT(*) |
| 618 | + FROM images |
| 619 | + WHERE images.starred = TRUE AND (1=1{query_conditions}) |
| 620 | + """ |
| 621 | + cursor.execute(starred_count_query, query_params) |
| 622 | + starred_count = cast(int, cursor.fetchone()[0]) |
| 623 | + |
| 624 | + # Get all image names with proper ordering |
| 625 | + if starred_first: |
| 626 | + names_query = f"""--sql |
| 627 | + SELECT images.image_name |
| 628 | + FROM images |
| 629 | + WHERE 1=1{query_conditions} |
| 630 | + ORDER BY images.starred DESC, images.created_at {order_dir.value} |
| 631 | + """ |
| 632 | + else: |
| 633 | + names_query = f"""--sql |
| 634 | + SELECT images.image_name |
| 635 | + FROM images |
| 636 | + WHERE 1=1{query_conditions} |
| 637 | + ORDER BY images.created_at {order_dir.value} |
| 638 | + """ |
| 639 | + |
| 640 | + cursor.execute(names_query, query_params) |
| 641 | + result = cast(list[sqlite3.Row], cursor.fetchall()) |
| 642 | + image_names = [row[0] for row in result] |
| 643 | + |
| 644 | + return ImageNamesResult(image_names=image_names, starred_count=starred_count, total_count=len(image_names)) |
0 commit comments