refactor: unify LoRA archive query skeleton
This commit is contained in:
@@ -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())
|
||||
|
||||
@@ -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:")
|
||||
|
||||
Reference in New Issue
Block a user