diff --git a/backend/app/database/images.py b/backend/app/database/images.py index ec9541a56..fd22af82b 100644 --- a/backend/app/database/images.py +++ b/backend/app/database/images.py @@ -6,6 +6,7 @@ from app.config.settings import ( DATABASE_PATH, ) +from app.database.connection import get_db_connection from app.logging.setup_logging import get_logger # Initialize logger @@ -42,48 +43,38 @@ class UntaggedImageRecord(TypedDict): ImageClassPair = Tuple[ImageId, ClassId] -def _connect() -> sqlite3.Connection: - conn = sqlite3.connect(DATABASE_PATH) - # Ensure ON DELETE CASCADE and other FKs are enforced - conn.execute("PRAGMA foreign_keys = ON") - return conn - - def db_create_images_table() -> None: - conn = _connect() - cursor = conn.cursor() + with get_db_connection() as conn: + cursor = conn.cursor() - # Create new images table with merged fields - cursor.execute( + # Create new images table with merged fields + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS images ( + id TEXT PRIMARY KEY, + path VARCHAR UNIQUE, + folder_id INTEGER, + thumbnailPath TEXT UNIQUE, + metadata TEXT, + isTagged BOOLEAN DEFAULT 0, + isFavourite BOOLEAN DEFAULT 0, + FOREIGN KEY (folder_id) REFERENCES folders(folder_id) ON DELETE CASCADE + ) """ - CREATE TABLE IF NOT EXISTS images ( - id TEXT PRIMARY KEY, - path VARCHAR UNIQUE, - folder_id INTEGER, - thumbnailPath TEXT UNIQUE, - metadata TEXT, - isTagged BOOLEAN DEFAULT 0, - isFavourite BOOLEAN DEFAULT 0, - FOREIGN KEY (folder_id) REFERENCES folders(folder_id) ON DELETE CASCADE ) - """ - ) - # Create new image_classes junction table - cursor.execute( + # Create new image_classes junction table + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS image_classes ( + image_id TEXT, + class_id INTEGER, + PRIMARY KEY (image_id, class_id), + FOREIGN KEY (image_id) REFERENCES images(id) ON DELETE CASCADE, + FOREIGN KEY (class_id) REFERENCES mappings(class_id) ON DELETE CASCADE + ) """ - CREATE TABLE IF NOT EXISTS image_classes ( - image_id TEXT, - class_id INTEGER, - PRIMARY KEY (image_id, class_id), - FOREIGN KEY (image_id) REFERENCES images(id) ON DELETE CASCADE, - FOREIGN KEY (class_id) REFERENCES mappings(class_id) ON DELETE CASCADE ) - """ - ) - - conn.commit() - conn.close() def db_bulk_insert_images(image_records: List[ImageRecord]) -> bool: @@ -91,33 +82,28 @@ def db_bulk_insert_images(image_records: List[ImageRecord]) -> bool: if not image_records: return True - conn = _connect() - cursor = conn.cursor() - try: - cursor.executemany( - """ - INSERT INTO images (id, path, folder_id, thumbnailPath, metadata, isTagged) - VALUES (:id, :path, :folder_id, :thumbnailPath, :metadata, :isTagged) - ON CONFLICT(path) DO UPDATE SET - folder_id=excluded.folder_id, - thumbnailPath=excluded.thumbnailPath, - metadata=excluded.metadata, - isTagged=CASE - WHEN excluded.isTagged THEN 1 - ELSE images.isTagged - END - """, - image_records, - ) - conn.commit() - return True + with get_db_connection() as conn: + cursor = conn.cursor() + cursor.executemany( + """ + INSERT INTO images (id, path, folder_id, thumbnailPath, metadata, isTagged) + VALUES (:id, :path, :folder_id, :thumbnailPath, :metadata, :isTagged) + ON CONFLICT(path) DO UPDATE SET + folder_id=excluded.folder_id, + thumbnailPath=excluded.thumbnailPath, + metadata=excluded.metadata, + isTagged=CASE + WHEN excluded.isTagged THEN 1 + ELSE images.isTagged + END + """, + image_records, + ) + return True except Exception as e: logger.error(f"Error inserting image records: {e}") - conn.rollback() return False - finally: - conn.close() def db_get_all_images(tagged: Union[bool, None] = None) -> List[dict]: @@ -131,87 +117,85 @@ def db_get_all_images(tagged: Union[bool, None] = None) -> List[dict]: Returns: List of dictionaries containing all image data including tags """ - conn = _connect() - cursor = conn.cursor() - try: - # Build the query with optional WHERE clause - query = """ - SELECT - i.id, - i.path, - i.folder_id, - i.thumbnailPath, - i.metadata, - i.isTagged, - i.isFavourite, - m.name as tag_name - FROM images i - LEFT JOIN image_classes ic ON i.id = ic.image_id - LEFT JOIN mappings m ON ic.class_id = m.class_id - """ - - params = [] - if tagged is not None: - query += " WHERE i.isTagged = ?" - params.append(tagged) - - query += " ORDER BY i.path, m.name" - - cursor.execute(query, params) - - results = cursor.fetchall() - - # Group results by image ID - images_dict = {} - for ( - image_id, - path, - folder_id, - thumbnail_path, - metadata, - is_tagged, - is_favourite, - tag_name, - ) in results: - if image_id not in images_dict: - # Safely parse metadata JSON -> dict - from app.utils.images import image_util_parse_metadata - - metadata_dict = image_util_parse_metadata(metadata) - - images_dict[image_id] = { - "id": image_id, - "path": path, - "folder_id": str(folder_id), - "thumbnailPath": thumbnail_path, - "metadata": metadata_dict, - "isTagged": bool(is_tagged), - "isFavourite": bool(is_favourite), - "tags": [], - } - - # Add tag if it exists (avoid duplicates) - if tag_name and tag_name not in images_dict[image_id]["tags"]: - images_dict[image_id]["tags"].append(tag_name) - - # Convert to list and set tags to None if empty - images = [] - for image_data in images_dict.values(): - if not image_data["tags"]: - image_data["tags"] = None - images.append(image_data) - - # Sort by path - images.sort(key=lambda x: x["path"]) + with get_db_connection() as conn: + cursor = conn.cursor() + + # Build the query with optional WHERE clause + query = """ + SELECT + i.id, + i.path, + i.folder_id, + i.thumbnailPath, + i.metadata, + i.isTagged, + i.isFavourite, + m.name as tag_name + FROM images i + LEFT JOIN image_classes ic ON i.id = ic.image_id + LEFT JOIN mappings m ON ic.class_id = m.class_id + """ - return images + params = [] + if tagged is not None: + query += " WHERE i.isTagged = ?" + params.append(tagged) + + query += " ORDER BY i.path, m.name" + + cursor.execute(query, params) + + results = cursor.fetchall() + + # Group results by image ID + images_dict = {} + for ( + image_id, + path, + folder_id, + thumbnail_path, + metadata, + is_tagged, + is_favourite, + tag_name, + ) in results: + if image_id not in images_dict: + # Safely parse metadata JSON -> dict + from app.utils.images import image_util_parse_metadata + + metadata_dict = image_util_parse_metadata(metadata) + + images_dict[image_id] = { + "id": image_id, + "path": path, + "folder_id": str(folder_id), + "thumbnailPath": thumbnail_path, + "metadata": metadata_dict, + "isTagged": bool(is_tagged), + "isFavourite": bool(is_favourite), + "tags": [], + } + + # Add tag if it exists (avoid duplicates) + if tag_name and tag_name not in images_dict[image_id]["tags"]: + images_dict[image_id]["tags"].append(tag_name) + + # Convert to list and set tags to None if empty + images = [] + for image_data in images_dict.values(): + if not image_data["tags"]: + image_data["tags"] = None + images.append(image_data) + + # Sort by path + images.sort(key=lambda x: x["path"]) + + return images except Exception as e: logger.error(f"Error getting all images: {e}") return [] - finally: - conn.close() def db_get_untagged_images() -> List[UntaggedImageRecord]: @@ -224,10 +208,9 @@ def db_get_untagged_images() -> List[UntaggedImageRecord]: Returns: List of dictionaries containing image data: id, path, folder_id, thumbnailPath, metadata """ - conn = _connect() - cursor = conn.cursor() + with get_db_connection() as conn: + cursor = conn.cursor() - try: cursor.execute( """ SELECT i.id, i.path, i.folder_id, i.thumbnailPath, i.metadata @@ -257,9 +240,6 @@ def db_get_untagged_images() -> List[UntaggedImageRecord]: return untagged_images - finally: - conn.close() - def db_update_image_tagged_status(image_id: ImageId, is_tagged: bool = True) -> bool: """ @@ -272,22 +252,17 @@ def db_update_image_tagged_status(image_id: ImageId, is_tagged: bool = True) -> Returns: True if update was successful, False otherwise """ - conn = _connect() - cursor = conn.cursor() - try: - cursor.execute( - "UPDATE images SET isTagged = ? WHERE id = ?", - (is_tagged, image_id), - ) - conn.commit() - return cursor.rowcount > 0 + with get_db_connection() as conn: + cursor = conn.cursor() + cursor.execute( + "UPDATE images SET isTagged = ? WHERE id = ?", + (is_tagged, image_id), + ) + return cursor.rowcount > 0 except Exception as e: logger.error(f"Error updating image tagged status: {e}") - conn.rollback() return False - finally: - conn.close() def db_insert_image_classes_batch(image_class_pairs: List[ImageClassPair]) -> bool: @@ -303,25 +278,20 @@ def db_insert_image_classes_batch(image_class_pairs: List[ImageClassPair]) -> bo if not image_class_pairs: return True - conn = _connect() - cursor = conn.cursor() - try: - cursor.executemany( - """ - INSERT OR IGNORE INTO image_classes (image_id, class_id) - VALUES (?, ?) - """, - image_class_pairs, - ) - conn.commit() - return True + with get_db_connection() as conn: + cursor = conn.cursor() + cursor.executemany( + """ + INSERT OR IGNORE INTO image_classes (image_id, class_id) + VALUES (?, ?) + """, + image_class_pairs, + ) + return True except Exception as e: logger.error(f"Error inserting image classes: {e}") - conn.rollback() return False - finally: - conn.close() def db_get_images_by_folder_ids( @@ -339,26 +309,24 @@ def db_get_images_by_folder_ids( if not folder_ids: return [] - conn = _connect() - cursor = conn.cursor() - try: - # Create placeholders for the IN clause - placeholders = ",".join("?" for _ in folder_ids) - cursor.execute( - f""" - SELECT id, path, thumbnailPath - FROM images - WHERE folder_id IN ({placeholders}) - """, - folder_ids, - ) - return cursor.fetchall() + with get_db_connection() as conn: + cursor = conn.cursor() + + # Create placeholders for the IN clause + placeholders = ",".join("?" for _ in folder_ids) + cursor.execute( + f""" + SELECT id, path, thumbnailPath + FROM images + WHERE folder_id IN ({placeholders}) + """, + folder_ids, + ) + return cursor.fetchall() except Exception as e: logger.error(f"Error getting images by folder IDs: {e}") return [] - finally: - conn.close() def db_delete_images_by_ids(image_ids: List[ImageId]) -> bool: @@ -375,47 +343,39 @@ def db_delete_images_by_ids(image_ids: List[ImageId]) -> bool: if not image_ids: return True - conn = _connect() - cursor = conn.cursor() - try: - # Create placeholders for the IN clause - placeholders = ",".join("?" for _ in image_ids) - cursor.execute( - f"DELETE FROM images WHERE id IN ({placeholders})", - image_ids, - ) - conn.commit() - logger.info(f"Deleted {cursor.rowcount} obsolete image(s) from database") - return True + with get_db_connection() as conn: + cursor = conn.cursor() + + # Create placeholders for the IN clause + placeholders = ",".join("?" for _ in image_ids) + cursor.execute( + f"DELETE FROM images WHERE id IN ({placeholders})", + image_ids, + ) + logger.info(f"Deleted {cursor.rowcount} obsolete image(s) from database") + return True except Exception as e: logger.error(f"Error deleting images: {e}") - conn.rollback() return False - finally: - conn.close() def db_toggle_image_favourite_status(image_id: str) -> bool: - conn = sqlite3.connect(DATABASE_PATH) - cursor = conn.cursor() try: - cursor.execute("SELECT id FROM images WHERE id = ?", (image_id,)) - if not cursor.fetchone(): - return False - cursor.execute( - """ - UPDATE images - SET isFavourite = CASE WHEN isFavourite = 1 THEN 0 ELSE 1 END - WHERE id = ? - """, - (image_id,), - ) - conn.commit() - return cursor.rowcount > 0 + with get_db_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT id FROM images WHERE id = ?", (image_id,)) + if not cursor.fetchone(): + return False + cursor.execute( + """ + UPDATE images + SET isFavourite = CASE WHEN isFavourite = 1 THEN 0 ELSE 1 END + WHERE id = ? + """, + (image_id,), + ) + return cursor.rowcount > 0 except Exception as e: logger.error(f"Database error: {e}") - conn.rollback() return False - finally: - conn.close()