nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: page image listings in SQL and drop three unused helpers

Browse Source 
get_recent_images and search_images_by_prompt take limit/offset so the recent
and search image routes no longer over-fetch a 5000-row window and slice it.
generate_content_hash, PromptModel.migrate_database and find_duplicates had no
callers outside their own tests. The isolated logger tests re-apply their
thread filter after update_config rebuilds the handlers. Long lines in the
touched database and image modules are wrapped so the full flake8 hook passes.
This commit is contained in:
Vito Sansevero
2026-10-08 18:01:34 -07:00
parent 11aea013ca
commit 0393d01746
9 changed files with 180 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.flake8
+1 -1
View File
@@ -1,6 +1,6 @@
[flake8]
max-line-length = 88
extend-ignore = E203, W503
exclude = reference/, .git, __pycache__, web/, venv/, htmlcov/, build/, dist/
exclude = reference/, .git, __pycache__, web/, venv/, htmlcov/, build/, dist/, .claude/, docs/
per-file-ignores =
tests/*:E402
database/models.py
+34 -26
View File
@@ -28,7 +28,10 @@ ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00"
def utc_now_iso() -> str:
"""Current UTC time as ISO text with microseconds, e.g. 2026-10-07T18:39:06.758123+00:00."""
"""Current UTC time as ISO text with microseconds.
Example: 2026-10-07T18:39:06.758123+00:00
"""
return datetime.datetime.now(datetime.timezone.utc).isoformat(
timespec="microseconds"
)
@@ -242,10 +245,13 @@ class PromptModel:
"CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)",
"CREATE INDEX IF NOT EXISTS idx_prompts_last_used ON prompts(last_used_at)",
"CREATE INDEX IF NOT EXISTS idx_prompts_run_count ON prompts(run_count)",
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)",
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON "
"generated_images(prompt_id)",
"CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)",
"CREATE INDEX IF NOT EXISTS idx_image_file_path ON generated_images(file_path)",
"CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)",
"CREATE INDEX IF NOT EXISTS idx_image_file_path ON "
"generated_images(file_path)",
"CREATE INDEX IF NOT EXISTS idx_generation_time ON "
"generated_images(generation_time)",
"CREATE INDEX IF NOT EXISTS idx_prompt_tags_tag ON prompt_tags(tag_id)",
"CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name)",
]
@@ -353,8 +359,10 @@ class PromptModel:
)
""")
conn.execute("""
INSERT INTO prompts_new (id, text, created_at, updated_at, category, tags, rating, notes, hash)
SELECT id, text, created_at, updated_at, category, tags, rating, notes, hash
INSERT INTO prompts_new (id, text, created_at, updated_at,
category, tags, rating, notes, hash)
SELECT id, text, created_at, updated_at, category, tags, rating,
notes, hash
FROM prompts
""")
conn.execute("DROP TABLE prompts")
@@ -408,8 +416,10 @@ class PromptModel:
INSERT INTO generated_images_new
(id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters)
SELECT id, CAST(prompt_id AS INTEGER), image_path, filename, generation_time,
file_size, width, height, format, workflow_data, prompt_metadata, parameters
SELECT id, CAST(prompt_id AS INTEGER), image_path, filename,
generation_time,
file_size, width, height, format, workflow_data,
prompt_metadata, parameters
FROM generated_images
WHERE prompt_id != '' AND prompt_id IS NOT NULL
AND CAST(prompt_id AS INTEGER) IN (SELECT id FROM prompts)
@@ -452,7 +462,8 @@ class PromptModel:
conn.execute(
"INSERT OR IGNORE INTO tags (name) "
"SELECT DISTINCT je.value FROM prompts, json_each(prompts.tags) AS je "
"WHERE prompts.tags IS NOT NULL AND prompts.tags != '' AND prompts.tags != '[]'"
"WHERE prompts.tags IS NOT NULL AND prompts.tags != '' AND "
"prompts.tags != '[]'"
)
# Populate junction table
@@ -467,7 +478,8 @@ class PromptModel:
tag_count = conn.execute("SELECT COUNT(*) FROM tags").fetchone()[0]
link_count = conn.execute("SELECT COUNT(*) FROM prompt_tags").fetchone()[0]
self.logger.info(
f"Tag migration complete: {tag_count} unique tags, {link_count} prompt-tag links"
f"Tag migration complete: {tag_count} unique tags, {link_count} "
"prompt-tag links"
)
except sqlite3.Error as e:
@@ -498,7 +510,8 @@ class PromptModel:
if "run_count" not in columns:
self.logger.info("Adding prompt usage column run_count")
conn.execute(
"ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL DEFAULT 0"
"ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL "
"DEFAULT 0"
)
added = True
if not added:
@@ -569,7 +582,8 @@ class PromptModel:
prompt_metadata TEXT,
parameters TEXT,
file_path TEXT,
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE,
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE
CASCADE,
UNIQUE(prompt_id, file_path)
)
""")
@@ -579,8 +593,10 @@ class PromptModel:
(id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters,
file_path)
SELECT id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters,
SELECT id, prompt_id, image_path, filename, generation_time,
file_size,
width, height, format, workflow_data, prompt_metadata,
parameters,
?
FROM generated_images WHERE id = ?
""",
@@ -599,16 +615,6 @@ class PromptModel:
self.logger.error(f"Migration error (image path uniqueness): {e}")
raise
def migrate_database(self) -> None:
"""
Apply any pending database migrations.
This method serves as an entry point for future schema migrations.
Add new migration logic here as the database evolves.
"""
# Future migrations can be added here
pass
def vacuum_database(self) -> None:
"""
Optimize database by running VACUUM command.
@@ -636,12 +642,14 @@ class PromptModel:
total_prompts = cursor.fetchone()["total_prompts"]
cursor = conn.execute(
"SELECT COUNT(DISTINCT category) as unique_categories FROM prompts WHERE category IS NOT NULL"
"SELECT COUNT(DISTINCT category) as unique_categories FROM "
"prompts WHERE category IS NOT NULL"
)
unique_categories = cursor.fetchone()["unique_categories"]
cursor = conn.execute(
"SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS NOT NULL"
"SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS "
"NOT NULL"
)
avg_rating = cursor.fetchone()["avg_rating"]
database/operations.py
+47 -94
View File
@@ -154,7 +154,8 @@ class PromptDatabase:
raise ValueError("Rating must be between 1 and 5")
self.logger.debug(
f"Saving prompt: text_length={len(text)}, category={category}, tags={tags}, rating={rating}"
f"Saving prompt: text_length={len(text)}, category={category}, "
f"tags={tags}, rating={rating}"
)
try:
@@ -424,7 +425,8 @@ class PromptDatabase:
"""
with self.model.get_connection() as conn:
cursor = conn.execute(
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE category = ? ORDER BY created_at DESC LIMIT ?",
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE category = ? "
"ORDER BY created_at DESC LIMIT ?",
(category, limit),
)
rows = cursor.fetchall()
@@ -574,7 +576,8 @@ class PromptDatabase:
"""
with self.model.get_connection() as conn:
cursor = conn.execute(
"SELECT DISTINCT TRIM(category) as category FROM prompts WHERE category IS NOT NULL AND TRIM(category) != '' ORDER BY category"
"SELECT DISTINCT TRIM(category) as category FROM prompts WHERE "
"category IS NOT NULL AND TRIM(category) != '' ORDER BY category"
)
return [row["category"] for row in cursor.fetchall()]
@@ -949,7 +952,8 @@ class PromptDatabase:
conn.commit()
self.logger.info(
f"Merged {tags_merged} tags into '{target_tag}', affected {affected} prompts"
f"Merged {tags_merged} tags into '{target_tag}', affected {affected} "
"prompts"
)
return {
"success": True,
@@ -968,7 +972,8 @@ class PromptDatabase:
with self.model.get_connection() as conn:
cursor = conn.execute(
"SELECT COUNT(*) as total FROM prompts "
"WHERE NOT EXISTS (SELECT 1 FROM prompt_tags WHERE prompt_id = prompts.id)"
"WHERE NOT EXISTS (SELECT 1 FROM prompt_tags WHERE prompt_id = "
"prompts.id)"
)
return cursor.fetchone()["total"]
@@ -1081,7 +1086,8 @@ class PromptDatabase:
tag_id = tag_map.get(tag_name)
if tag_id:
conn.execute(
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) VALUES (?, ?)",
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) "
"VALUES (?, ?)",
(prompt_id, tag_id),
)
@@ -1170,79 +1176,6 @@ class PromptDatabase:
self.logger.error(f"Error exporting prompts: {e}")
return False
def find_duplicates(self) -> List[Dict[str, Any]]:
"""
Find duplicate prompts based on text content without removing them.
Returns:
List of duplicate groups, each containing:
- text: The duplicate text content
- prompts: List of prompt records with same text
"""
self.logger.info("Scanning for duplicate prompts")
try:
with self.model.get_connection() as conn:
# Find duplicates by text content (case-insensitive)
# Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility
# We'll sort the IDs manually after fetching
cursor = conn.execute("""
SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count,
GROUP_CONCAT(id) as ids,
GROUP_CONCAT(created_at) as created_dates
FROM prompts
GROUP BY LOWER(TRIM(text))
HAVING COUNT(*) > 1
""")
duplicate_groups = cursor.fetchall()
self.logger.debug(
f"Found {len(duplicate_groups)} groups of duplicate prompts"
)
result = []
for group in duplicate_groups:
ids = group["ids"].split(",")
created_dates = group["created_dates"].split(",")
# Sort IDs by created_at date
id_date_pairs = list(zip(ids, created_dates))
id_date_pairs.sort(key=lambda x: x[1]) # Sort by date
ids = [pair[0] for pair in id_date_pairs]
# Get full details for all prompts in this duplicate group
prompts = []
for prompt_id in ids:
cursor = conn.execute(
"SELECT id, text, category, rating, created_at, updated_at "
"FROM prompts WHERE id = ?",
(int(prompt_id),),
)
prompt_data = cursor.fetchone()
if prompt_data:
prompt_dict = dict(prompt_data)
prompt_dict["tags"] = self._get_prompt_tags(
conn, int(prompt_id)
)
prompts.append(prompt_dict)
if prompts:
result.append(
{
"text": prompts[0][
"text"
], # Use the actual text (not normalized)
"prompts": prompts,
}
)
self.logger.info(f"Found {len(result)} groups with duplicates")
return result
except Exception as e:
self.logger.error(f"Error finding duplicates: {e}")
return []
def cleanup_duplicates(self) -> int:
"""
Remove duplicate prompts based on text content, preserving all image links.
@@ -1287,7 +1220,8 @@ class PromptDatabase:
duplicate_ids = sorted_ids[1:]
self.logger.debug(
f"Merging duplicates: keeping {primary_id}, removing {duplicate_ids}"
f"Merging duplicates: keeping {primary_id}, removing "
f"{duplicate_ids}"
)
# Get primary prompt details
@@ -1327,7 +1261,8 @@ class PromptDatabase:
if total_removed > 0:
self.logger.info(
f"Removed {total_removed} duplicate prompts, transferred {total_images_transferred} images"
f"Removed {total_removed} duplicate prompts, transferred "
f"{total_images_transferred} images"
)
else:
self.logger.info("No duplicate prompts found")
@@ -1380,7 +1315,8 @@ class PromptDatabase:
)
if not cursor.fetchone():
self.logger.warning(
f"Prompt ID {prompt_id_int} not found in database, skipping image linking"
f"Prompt ID {prompt_id_int} not found in database, skipping "
"image linking"
)
return 0
@@ -1455,12 +1391,15 @@ class PromptDatabase:
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
def get_recent_images(self, limit: int = 50) -> List[Dict[str, Any]]:
def get_recent_images(
self, limit: int = 50, offset: int = 0
) -> List[Dict[str, Any]]:
"""
Get recently generated images across all prompts.
Get recently generated images across all prompts, newest first.
Args:
limit: Maximum number of images to return
offset: Number of newest images to skip (paging)
Returns:
List of image records with prompt text
@@ -1472,9 +1411,9 @@ class PromptDatabase:
FROM generated_images gi
LEFT JOIN prompts p ON gi.prompt_id = p.id
ORDER BY gi.generation_time DESC
LIMIT ?
LIMIT ? OFFSET ?
""",
(limit,),
(int(limit), max(0, int(offset))),
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
@@ -1496,7 +1435,8 @@ class PromptDatabase:
sql = (
"SELECT gi.*, p.text as prompt_text, "
"(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt "
"JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = p.id) AS _prompt_tags_list "
"JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = p.id) AS "
"_prompt_tags_list "
"FROM generated_images gi "
"INNER JOIN prompts p ON gi.prompt_id = p.id "
"WHERE gi.image_path IS NOT NULL AND gi.image_path != '' "
@@ -1520,12 +1460,16 @@ class PromptDatabase:
result.append(data)
return result
def search_images_by_prompt(self, search_term: str) -> List[Dict[str, Any]]:
def search_images_by_prompt(
self, search_term: str, limit: Optional[int] = None, offset: int = 0
) -> List[Dict[str, Any]]:
"""
Search images by prompt text.
Search images by prompt text, newest first.
Args:
search_term: Text to search for in prompt content
limit: Maximum number of images to return (None for no limit)
offset: Number of newest matches to skip (paging)
Returns:
List of image records with prompt text
@@ -1538,8 +1482,13 @@ class PromptDatabase:
JOIN prompts p ON gi.prompt_id = p.id
WHERE p.text LIKE ? ESCAPE '\\'
ORDER BY gi.generation_time DESC
LIMIT ? OFFSET ?
""",
(f"%{escape_like(search_term)}%",),
(
f"%{escape_like(search_term)}%",
-1 if limit is None else int(limit),
max(0, int(offset)),
),
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
@@ -1710,7 +1659,8 @@ class PromptDatabase:
if cursor.rowcount > 0:
self.logger.debug(
f"Transferred {cursor.rowcount} images from prompt {dup_id} to {primary_id}"
f"Transferred {cursor.rowcount} images from prompt {dup_id} "
f"to {primary_id}"
)
return transferred_count
@@ -1888,7 +1838,8 @@ class PromptDatabase:
tag_id = tag_map.get(tag_name)
if tag_id:
cursor = conn.execute(
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) VALUES (?, ?)",
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) "
"VALUES (?, ?)",
(pid, tag_id),
)
if cursor.rowcount > 0:
@@ -2025,7 +1976,8 @@ class PromptDatabase:
""")
for ref in cursor.fetchall():
issues.append(
f"prompt_tags entry references non-existent prompt {ref['prompt_id']}"
"prompt_tags entry references non-existent prompt "
f"{ref['prompt_id']}"
)
# Check for orphaned image entries
@@ -2036,7 +1988,8 @@ class PromptDatabase:
""")
for ref in cursor.fetchall():
issues.append(
f"Image {ref['id']} references non-existent prompt {ref['prompt_id']}"
f"Image {ref['id']} references non-existent prompt "
f"{ref['prompt_id']}"
)
return issues
py/api/images.py
+31 -15
View File
@@ -149,7 +149,6 @@ def _thumbnail_complete_payload(total, stats, start):
# Caps on request-driven work (list sizes, scans, request bodies).
MAX_GALLERY_FILES = 50_000
MAX_RECENT_IMAGES_WINDOW = 5_000
MAX_BULK_LIMIT = 5_000
MAX_JSON_BODY_BYTES = 1024 * 1024
_READ_CHUNK = 64 * 1024
@@ -367,13 +366,9 @@ class ImageRoutesMixin:
except ValueError:
return bad_request("limit and offset must be integers")
# The DB query has no offset: fetch one bounded window and slice,
# so the offset is capped to keep the window small.
offset = min(offset, MAX_RECENT_IMAGES_WINDOW - limit)
rows = await self._run_in_executor(
self.db.get_recent_images, limit + offset
images = await self._run_in_executor(
self.db.get_recent_images, limit, offset
)
images = rows[offset:]
return web.json_response(
{
@@ -436,14 +431,32 @@ class ImageRoutesMixin:
{"success": False, "error": "Search query required"}, status=400
)
images = await self._run_in_executor(self.db.search_images_by_prompt, query)
try:
limit, offset = parse_page_params(request.query)
except ValueError:
return bad_request("limit and offset must be integers")
images = await self._run_in_executor(
self.db.search_images_by_prompt, query, limit, offset
)
return web.json_response(
{"success": True, "images": images, "query": query}
{
"success": True,
"images": images,
"query": query,
"pagination": {
"limit": limit,
"offset": offset,
"count": len(images),
},
}
)
except Exception as e:
self.logger.error(f"Search images error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": _safe_error(e)}, status=500
)
def _scan_gallery_files_sync(self, output_path):
"""Scan output directory for media files (blocking I/O, run in executor).
@@ -972,7 +985,6 @@ class ImageRoutesMixin:
async def clear_thumbnails(self, request):
"""Safely clear only our generated thumbnails, never touch original images."""
try:
import shutil
# Find ComfyUI output directory
output_dir = self._find_comfyui_output_dir()
@@ -1006,12 +1018,14 @@ class ImageRoutesMixin:
or thumbnails_dir.name != "thumbnails"
):
self.logger.error(
f"Safety check failed: thumbnails directory path invalid: {thumbnails_dir}"
"Safety check failed: thumbnails directory path invalid: "
f"{thumbnails_dir}"
)
return web.json_response(
{
"success": False,
"error": "Safety check failed: invalid thumbnails directory path",
"error": "Safety check failed: invalid thumbnails "
"directory path",
},
status=400,
)
@@ -1064,7 +1078,8 @@ class ImageRoutesMixin:
return f"{bytes_size:.1f} TB"
self.logger.info(
f"Thumbnail cleanup: cleared {cleared_count} files ({format_size(cleared_size)})"
f"Thumbnail cleanup: cleared {cleared_count} files "
f"({format_size(cleared_size)})"
)
return web.json_response(
@@ -1073,7 +1088,8 @@ class ImageRoutesMixin:
"cleared_files": cleared_count,
"cleared_size": cleared_size,
"cleared_size_formatted": format_size(cleared_size),
"message": f"Cleared {cleared_count} thumbnail files ({format_size(cleared_size)})",
"message": f"Cleared {cleared_count} thumbnail files "
"({format_size(cleared_size)})",
}
)
tests/test_api_images.py
+8 -9
View File
@@ -502,9 +502,9 @@ class TestRecentImagesCaps(ImageAPITestCase):
seen = []
original = self.api.db.get_recent_images
def spy(limit):
seen.append(limit)
return original(limit)
def spy(limit, offset=0):
seen.append((limit, offset))
return original(limit, offset)
self.api.db.get_recent_images = spy
return seen
@@ -516,7 +516,7 @@ class TestRecentImagesCaps(ImageAPITestCase):
data = await resp.json()
self.assertEqual(data["pagination"]["limit"], 50)
self.assertEqual(seen, [50])
self.assertEqual(seen, [(50, 0)])
async def test_huge_offset_is_clamped_before_reaching_db(self):
seen = self._spy_db()
@@ -527,10 +527,9 @@ class TestRecentImagesCaps(ImageAPITestCase):
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertLessEqual(seen[0], images_module.MAX_RECENT_IMAGES_WINDOW)
self.assertEqual(
data["pagination"]["offset"], images_module.MAX_RECENT_IMAGES_WINDOW - 10
)
# The offset is paged in SQL, so the DB sees the clamped value itself.
self.assertEqual(seen[0], (10, data["pagination"]["offset"]))
self.assertLess(data["pagination"]["offset"], 999999999999)
async def test_float_string_limit_is_400(self):
resp = await self.client.request(
@@ -544,7 +543,7 @@ class TestRecentImagesCaps(ImageAPITestCase):
await self.client.request("GET", "/prompt_manager/images/recent?limit=0")
self.assertEqual(seen, [1])
self.assertEqual(seen, [(1, 0)])
class TestAllImagesCaps(ImageAPITestCase):
tests/test_database.py
+37 -15
View File
@@ -646,18 +646,6 @@ class TestDuplicateMerging(DatabaseTestCase):
)
return pid
def test_find_duplicates_groups_case_insensitively_oldest_first(self):
newer = self._save_dup("Same Text", "2026-02-01T00:00:00", tags=["t2"])
older = self._save_dup("same text ", "2026-01-01T00:00:00", tags=["t1"])
self._save_dup("unique", "2026-01-01T00:00:00")
groups = self.db.find_duplicates()
self.assertEqual(len(groups), 1)
self.assertEqual([p["id"] for p in groups[0]["prompts"]], [older, newer])
self.assertEqual(groups[0]["text"], "same text")
self.assertEqual(groups[0]["prompts"][0]["tags"], ["t1"])
def test_cleanup_keeps_oldest_merges_metadata_and_moves_images(self):
keep = self._save_dup(
"dup", "2026-01-01T00:00:00", tags=["a"], notes="first", rating=2
@@ -694,7 +682,6 @@ class TestDuplicateMerging(DatabaseTestCase):
with patch.object(
self.db.model, "get_connection", side_effect=sqlite3.OperationalError
):
self.assertEqual(self.db.find_duplicates(), [])
self.assertEqual(self.db.cleanup_duplicates(), 0)
def test_hash_duplicates_are_impossible_under_unique_hash(self):
@@ -1083,10 +1070,9 @@ class TestModelHousekeeping(DatabaseTestCase):
):
self.assertEqual(self.db.model.get_database_info(), {})
def test_vacuum_and_noop_migrate(self):
def test_vacuum_logs_errors_instead_of_raising(self):
self._save("x")
self.db.model.vacuum_database()
self.db.model.migrate_database()
with patch(
"database.models.sqlite3.connect", side_effect=sqlite3.OperationalError
):
@@ -1106,3 +1092,39 @@ class TestModelHousekeeping(DatabaseTestCase):
if __name__ == "__main__":
unittest.main()
class TestImagePaging(DatabaseTestCase):
"""Image listings page in the database, not by over-fetching in the API."""
def _link_many(self, count):
pid = self._save("paged prompt")
image_ids = [
self.db.link_image_to_prompt(pid, f"/out/img{i:03d}.png")
for i in range(count)
]
# Stamp distinct times afterwards through the model's own connection,
# so no second connection holds a write lock while linking.
with self.db.model.get_connection() as conn:
for i, image_id in enumerate(image_ids):
conn.execute(
"UPDATE generated_images SET generation_time = ? WHERE id = ?",
(f"2026-01-01T00:00:{i:02d}", image_id),
)
return pid
def test_recent_images_offset_skips_newest_rows(self):
self._link_many(5)
newest_first = [img["filename"] for img in self.db.get_recent_images(limit=5)]
page = self.db.get_recent_images(limit=2, offset=2)
self.assertEqual([img["filename"] for img in page], newest_first[2:4])
self.assertEqual(self.db.get_recent_images(limit=2, offset=10), [])
def test_search_images_by_prompt_is_bounded_and_pages(self):
self._link_many(5)
first = self.db.search_images_by_prompt("paged", limit=2)
second = self.db.search_images_by_prompt("paged", limit=2, offset=2)
self.assertEqual(len(first), 2)
self.assertEqual(len(second), 2)
self.assertNotEqual({img["id"] for img in first}, {img["id"] for img in second})
self.assertEqual(len(self.db.search_images_by_prompt("paged")), 5)
tests/test_hashing.py
+1 -76
View File
@@ -1,7 +1,7 @@
"""
Tests for hashing utilities.
Tests generate_content_hash and is_duplicate_prompt which are
Tests is_duplicate_prompt and generate_prompt_hash, which are
currently untested. generate_prompt_hash is covered in test_basic.py
but we add edge cases here.
"""
@@ -14,7 +14,6 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.hashing import (
generate_prompt_hash,
generate_content_hash,
is_duplicate_prompt,
)
@@ -63,80 +62,6 @@ class TestSingleHashImplementation(unittest.TestCase):
self.assertFalse(os.path.exists(os.path.join(ROOT, "restart_gallery.py")))
class TestGenerateContentHash(unittest.TestCase):
"""Test generate_content_hash function."""
def test_basic_content(self):
content = {
"text": "A prompt",
"category": "nature",
"tags": ["sky"],
"workflow_name": "wf1",
}
h = generate_content_hash(content)
self.assertEqual(len(h), 64) # SHA256 hex digest
def test_normalization_case_insensitive(self):
c1 = {"text": "Hello World", "category": "Test"}
c2 = {"text": "hello world", "category": "test"}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_normalization_strips_whitespace(self):
c1 = {"text": " hello ", "category": " test "}
c2 = {"text": "hello", "category": "test"}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_tag_order_independent(self):
c1 = {"text": "prompt", "tags": ["beta", "alpha"]}
c2 = {"text": "prompt", "tags": ["alpha", "beta"]}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_different_text_different_hash(self):
c1 = {"text": "prompt A"}
c2 = {"text": "prompt B"}
self.assertNotEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_different_category_different_hash(self):
c1 = {"text": "same", "category": "cat1"}
c2 = {"text": "same", "category": "cat2"}
self.assertNotEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_different_tags_different_hash(self):
c1 = {"text": "same", "tags": ["tag1"]}
c2 = {"text": "same", "tags": ["tag2"]}
self.assertNotEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_empty_dict(self):
h = generate_content_hash({})
self.assertEqual(len(h), 64)
def test_missing_fields_use_defaults(self):
c1 = {"text": "hello"}
c2 = {"text": "hello", "category": None, "tags": [], "workflow_name": None}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_none_category_same_as_empty(self):
c1 = {"text": "test", "category": None}
c2 = {"text": "test", "category": ""}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_none_workflow_same_as_empty(self):
c1 = {"text": "test", "workflow_name": None}
c2 = {"text": "test", "workflow_name": ""}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_empty_tags_filtered(self):
c1 = {"text": "test", "tags": ["valid", "", " "]}
c2 = {"text": "test", "tags": ["valid"]}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_deterministic(self):
content = {"text": "stable", "tags": ["a", "b"]}
h1 = generate_content_hash(content)
h2 = generate_content_hash(content)
self.assertEqual(h1, h2)
class TestIsDuplicatePrompt(unittest.TestCase):
"""Test is_duplicate_prompt function."""
tests/test_utils_logging_config.py
+21 -7
View File
@@ -35,15 +35,26 @@ class IsolatedLoggerTestCase(unittest.TestCase):
self.manager.update_config({"console_logging": False})
# Drop the "Updated logging configuration" entry.
self.manager._log_buffer.clear()
# Other tests leave daemon threads (image monitor, watchdog) logging under
# "prompt_manager.*"; only records from this test's thread may reach the buffer.
# The same goes for the rotating file handler: a stray line after a
# truncate would roll the file over again and resurrect the backups.
# Other tests leave daemon threads (image monitor, watchdog) logging
# under "prompt_manager.*"; only records from this test's thread may
# reach the buffer or the rotating file (a stray line after a truncate
# would roll the file over and resurrect the backups). update_config
# rebuilds the handlers, so the filter is re-applied after every call.
self._filter_to_this_thread()
original_update = self.manager.update_config
def update_and_refilter(new_config):
original_update(new_config)
self._filter_to_this_thread()
self.manager.update_config = update_and_refilter
self.logger = logging.getLogger("prompt_manager.coverage_test")
self.addCleanup(self._restore_global_logging)
def _filter_to_this_thread(self):
this_thread = threading.get_ident()
for handler in self.manager.logger.handlers:
handler.addFilter(lambda record: record.thread == this_thread)
self.logger = logging.getLogger("prompt_manager.coverage_test")
self.addCleanup(self._restore_global_logging)
def _restore_global_logging(self):
for handler in logging.getLogger("prompt_manager").handlers[:]:
@@ -186,7 +197,10 @@ class TestFilesAndRotation(IsolatedLoggerTestCase):
self.assertEqual(result["errors"], [])
self.logger.info("after truncate")
leftovers = sorted(p.name for p in pathlib.Path(self.log_dir).iterdir())
self.assertEqual(leftovers, ["prompt_manager.log"])
contents = {
name: pathlib.Path(self.log_dir, name).read_bytes() for name in leftovers
}
self.assertEqual(leftovers, ["prompt_manager.log"], contents)
main_log = pathlib.Path(self.log_dir, "prompt_manager.log").read_bytes()
self.assertNotIn(b"\x00", main_log)
self.assertIn(b"after truncate", main_log)
utils/hashing.py
-48
View File
@@ -65,54 +65,6 @@ def generate_prompt_hash(text: str) -> str:
return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest()
def generate_content_hash(content: dict) -> str:
"""
Generate a hash for prompt content including metadata.
Creates a comprehensive hash that includes not just the prompt text
but also associated metadata like category, tags, and workflow name.
This enables detection of prompts that are identical in all aspects.
Args:
content: Dictionary containing prompt data with optional keys:
- text: The prompt text
- category: Prompt category
- tags: List of tags
- workflow_name: Associated workflow name
Returns:
SHA256 hexadecimal digest of the normalized content structure
Note:
The hash is generated from a normalized JSON representation with
sorted keys and normalized text fields to ensure consistency.
"""
import json
# Create a normalized representation of the content
normalized = {
"text": content.get("text", "").strip().lower(),
"category": (
content.get("category", "").strip().lower()
if content.get("category")
else ""
),
"tags": sorted(
[tag.strip().lower() for tag in content.get("tags", []) if tag.strip()]
),
"workflow_name": (
content.get("workflow_name", "").strip().lower()
if content.get("workflow_name")
else ""
),
}
# Convert to JSON string for consistent hashing
content_str = json.dumps(normalized, sort_keys=True)
return hashlib.sha256(content_str.encode("utf-8")).hexdigest()
def is_duplicate_prompt(text1: str, text2: str, threshold: float = 0.95) -> bool:
"""
Check if two prompts are likely duplicates using hash comparison.