refactor: unify LoRA archive query skeleton

This commit is contained in:
Dariusz L
2026-08-06 23:09:13 +02:00
parent 069a588fb1
commit 497e6d0658
2 changed files with 69 additions and 70 deletions
+24 -70
View File
@@ -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())
+45
View File
@@ -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:")