diff --git a/core/sources/lora_manager_archive.py b/core/sources/lora_manager_archive.py index e909a57..0f8a1a7 100644 --- a/core/sources/lora_manager_archive.py +++ b/core/sources/lora_manager_archive.py @@ -44,6 +44,19 @@ DEFAULT_DB_RELATIVE_PATH = os.path.join("civitai", "civitai.sqlite") _search_cache: Dict[str, Any] = {} _db_path_cache: Optional[str] = None GENERIC_FILENAME_TOKENS = get_generic_filename_tokens() +_CANDIDATE_SELECT_SQL = """ +SELECT + m.id AS model_id, + m.name AS model_name, + m.type AS model_type, + v.id AS version_id, + v.name AS version_name, + v.base_model AS base_model, + v.position AS position +FROM models m +JOIN model_versions v ON v.model_id = m.id +WHERE {where_clause} +""" def clear_search_cache(): @@ -280,21 +293,9 @@ def _query_candidate_rows( if first_term and second_term and version_term: query_plans.append( ( - """ - SELECT - m.id AS model_id, - m.name AS model_name, - m.type AS model_type, - v.id AS version_id, - v.name AS version_name, - v.base_model AS base_model, - v.position AS position - FROM models m - JOIN model_versions v ON v.model_id = m.id - WHERE m.name LIKE ? COLLATE NOCASE - AND m.name LIKE ? COLLATE NOCASE - AND v.name LIKE ? COLLATE NOCASE - """, + "m.name LIKE ? COLLATE NOCASE\n" + " AND m.name LIKE ? COLLATE NOCASE\n" + " AND v.name LIKE ? COLLATE NOCASE", [f"%{first_term}%", f"%{second_term}%", f"%{version_term}%"], ) ) @@ -302,20 +303,8 @@ def _query_candidate_rows( if first_term and second_term: query_plans.append( ( - """ - SELECT - m.id AS model_id, - m.name AS model_name, - m.type AS model_type, - v.id AS version_id, - v.name AS version_name, - v.base_model AS base_model, - v.position AS position - FROM models m - JOIN model_versions v ON v.model_id = m.id - WHERE m.name LIKE ? COLLATE NOCASE - AND m.name LIKE ? COLLATE NOCASE - """, + "m.name LIKE ? COLLATE NOCASE\n" + " AND m.name LIKE ? COLLATE NOCASE", [f"%{first_term}%", f"%{second_term}%"], ) ) @@ -323,20 +312,8 @@ def _query_candidate_rows( if first_term and version_term: query_plans.append( ( - """ - SELECT - m.id AS model_id, - m.name AS model_name, - m.type AS model_type, - v.id AS version_id, - v.name AS version_name, - v.base_model AS base_model, - v.position AS position - FROM models m - JOIN model_versions v ON v.model_id = m.id - WHERE m.name LIKE ? COLLATE NOCASE - AND v.name LIKE ? COLLATE NOCASE - """, + "m.name LIKE ? COLLATE NOCASE\n" + " AND v.name LIKE ? COLLATE NOCASE", [f"%{first_term}%", f"%{version_term}%"], ) ) @@ -344,19 +321,7 @@ def _query_candidate_rows( if first_term: query_plans.append( ( - """ - SELECT - m.id AS model_id, - m.name AS model_name, - m.type AS model_type, - v.id AS version_id, - v.name AS version_name, - v.base_model AS base_model, - v.position AS position - FROM models m - JOIN model_versions v ON v.model_id = m.id - WHERE m.name LIKE ? COLLATE NOCASE - """, + "m.name LIKE ? COLLATE NOCASE", [f"%{first_term}%"], ) ) @@ -364,19 +329,7 @@ def _query_candidate_rows( if version_term: query_plans.append( ( - """ - SELECT - m.id AS model_id, - m.name AS model_name, - m.type AS model_type, - v.id AS version_id, - v.name AS version_name, - v.base_model AS base_model, - v.position AS position - FROM models m - JOIN model_versions v ON v.model_id = m.id - WHERE v.name LIKE ? COLLATE NOCASE - """, + "v.name LIKE ? COLLATE NOCASE", [f"%{version_term}%"], ) ) @@ -384,7 +337,8 @@ def _query_candidate_rows( seen_versions = set() candidates: List[sqlite3.Row] = [] - for sql, params in query_plans: + for where_clause, params in query_plans: + sql = _CANDIDATE_SELECT_SQL.format(where_clause=where_clause) if model_type: sql += " AND m.type = ?" params.append(model_type.upper()) diff --git a/tests/test_lora_manager_archive.py b/tests/test_lora_manager_archive.py index 2055943..e25c1d6 100644 --- a/tests/test_lora_manager_archive.py +++ b/tests/test_lora_manager_archive.py @@ -34,6 +34,51 @@ class LoraManagerArchiveTests(unittest.TestCase): self.assertIn("model", tokens_generic["content_terms"]) self.assertIn("v1", tokens_generic["version_terms"]) + def test_query_candidate_rows_preserves_all_filter_plans(self): + class EmptyQueryResult: + def fetchall(self): + return [] + + class RecordingConnection: + def __init__(self): + self.calls = [] + + def execute(self, sql, params): + self.calls.append((sql, tuple(params))) + return EmptyQueryResult() + + connection = RecordingConnection() + + candidates = _query_candidate_rows( + connection, + "aurora_style_v2.safetensors", + "lora", + 2, + ) + + self.assertEqual([], candidates) + self.assertEqual(5, len(connection.calls)) + self.assertTrue( + all("m.type = ?" in sql for sql, _params in connection.calls) + ) + self.assertTrue( + all( + "ORDER BY m.id ASC, COALESCE(v.position, 999999) ASC LIMIT ?" + in sql + for sql, _params in connection.calls + ) + ) + self.assertEqual( + [ + ("%aurora%", "%style%", "%v2%", "LORA", 50), + ("%aurora%", "%style%", "LORA", 50), + ("%aurora%", "%v2%", "LORA", 50), + ("%aurora%", "LORA", 50), + ("%v2%", "LORA", 50), + ], + [params for _sql, params in connection.calls], + ) + def test_database_operations_with_in_memory_sqlite(self): # Create an in-memory database conn = sqlite3.connect(":memory:")