Added test for sqlite-utils search, refs #192

This commit is contained in:
Simon Willison 2020-11-06 15:40:42 -08:00
commit 63e2bdf18d
4 changed files with 36 additions and 3 deletions

View file

@ -941,6 +941,7 @@ def query(
"Execute SQL query and return the results as JSON"
db = sqlite_utils.Database(path)
_load_extensions(db, load_extension)
db.register_fts4_bm25()
with db.conn:
cursor = db.execute(sql, dict(param))
if cursor.description is None:

View file

@ -1360,7 +1360,7 @@ class Table(Queryable):
rank_implementation = "[{}].rank".format(fts_table)
else:
self.db.register_fts4_bm25()
rank_implementation = "-rank_bm25(matchinfo([{}], 'pcnalx'))".format(
rank_implementation = "rank_bm25(matchinfo([{}], 'pcnalx'))".format(
fts_table
)
return sql.format(

View file

@ -1677,3 +1677,35 @@ def test_insert_encoding(tmpdir):
"longitude": None,
},
]
@pytest.mark.parametrize("fts", ["FTS4", "FTS5"])
@pytest.mark.parametrize(
"extra_arg,expected",
[
(
None,
'[{"rowid": 2, "id": 2, "title": "Title the second", "rank": -0.5108256237659907}]\n',
),
("--csv", "rowid,id,title,rank\n2,2,Title the second,-0.5108256237659907\n"),
],
)
def test_search(tmpdir, fts, extra_arg, expected):
db_path = str(tmpdir / "test.db")
db = Database(db_path)
db["articles"].insert_all(
[
{"id": 1, "title": "Title the first"},
{"id": 2, "title": "Title the second"},
{"id": 3, "title": "Title the third"},
],
pk="id",
)
db["articles"].enable_fts(["title"], fts_version=fts)
result = CliRunner().invoke(
cli.cli,
["search", db_path, "articles", "second"] + ([extra_arg] if extra_arg else []),
catch_exceptions=False,
)
assert result.exit_code == 0
assert result.output == expected

View file

@ -116,7 +116,7 @@ def test_populate_fts(fresh_db):
"text": "racoons are biting trash pandas",
"country": "USA",
"not_searchable": "bar",
"rank": 0.5108256237659907,
"rank": -0.5108256237659907,
}
] == rows
@ -137,7 +137,7 @@ def test_populate_fts_escape_table_names(fresh_db):
"text": "racoons are biting trash pandas",
"country": "USA",
"not_searchable": "bar",
"rank": 0.5108256237659907,
"rank": -0.5108256237659907,
}
] == list(table.search("usa"))