Refactored to table.search_sql() method, added --limit

This commit is contained in:
Simon Willison 2020-11-03 14:46:18 -08:00
commit de39e8db1e
4 changed files with 119 additions and 45 deletions

View file

@ -907,6 +907,8 @@ You can specify a subset of columns to be returned using the ``-c`` option one o
$ sqlite-utils search mydb.db documents searchterm -c title -c created $ sqlite-utils search mydb.db documents searchterm -c title -c created
By default all search results will be returned. You can use ``--limit 20`` to return just the first 20 results.
Use the ``--sql`` option to output the SQL that would be executed, rather than running the query:: Use the ``--sql`` option to output the SQL that would be executed, rather than running the query::
$ sqlite-utils search mydb.db documents searchterm --sql $ sqlite-utils search mydb.db documents searchterm --sql
@ -923,11 +925,9 @@ Use the ``--sql`` option to output the SQL that would be executed, rather than r
[original] [original]
join [documents_fts] on [original].rowid = [documents_fts].rowid join [documents_fts] on [original].rowid = [documents_fts].rowid
where where
[documents_fts] match :search [documents_fts] match :query
order by order by
rank desc rank desc
limit
20
.. _cli_vacuum: .. _cli_vacuum:

View file

@ -7,7 +7,6 @@ import hashlib
import pathlib import pathlib
import sqlite_utils import sqlite_utils
from sqlite_utils.db import AlterError from sqlite_utils.db import AlterError
import textwrap
import itertools import itertools
import json import json
import os import os
@ -979,6 +978,12 @@ def query(
@click.argument("q") @click.argument("q")
@click.option("-o", "--order", type=str, help="Order by ('column' or 'column desc')") @click.option("-o", "--order", type=str, help="Order by ('column' or 'column desc')")
@click.option("-c", "--column", type=str, multiple=True, help="Columns to return") @click.option("-c", "--column", type=str, multiple=True, help="Columns to return")
@click.option(
"--limit",
type=int,
default=20,
help="Number of rows to return, default 20, set to 0 for unlimited",
)
@click.option( @click.option(
"--sql", "show_sql", is_flag=True, help="Show SQL query that would be run" "--sql", "show_sql", is_flag=True, help="Show SQL query that would be run"
) )
@ -993,6 +998,7 @@ def search(
order, order,
show_sql, show_sql,
column, column,
limit,
nl, nl,
arrays, arrays,
csv, csv,
@ -1009,55 +1015,19 @@ def search(
table_obj = db[dbtable] table_obj = db[dbtable]
if not table_obj.exists(): if not table_obj.exists():
raise click.ClickException("Table '{}' does not exist".format(dbtable)) raise click.ClickException("Table '{}' does not exist".format(dbtable))
fts_table = table_obj.detect_fts() if not table_obj.detect_fts():
if not fts_table:
raise click.ClickException( raise click.ClickException(
"Table '{}' is not configured for full-text search".format(dbtable) "Table '{}' is not configured for full-text search".format(dbtable)
) )
# Pick names for table and rank column that don't clash
original = "original_" if dbtable == "original" else "original"
rank = "rank"
while rank in table_obj.columns_dict:
rank = rank + "_"
columns = "*"
if column: if column:
# Check they all exist # Check they all exist
table_columns = table_obj.columns_dict
for c in column: for c in column:
if c not in table_obj.columns_dict: if c not in table_columns:
raise click.ClickException( raise click.ClickException(
"Table '{}' has no column '{}".format(dbtable, c) "Table '{}' has no column '{}".format(dbtable, c)
) )
columns = ", ".join("[{}]".format(c) for c in column) sql = table_obj.search_sql(columns=column, order=order, limit=limit)
sql = textwrap.dedent(
"""
with {original} as (
select
rowid,
{columns}
from [{dbtable}]
)
select
{original}.*,
[{fts}].rank as {rank}
from
[{original}]
join [{fts}] on [{original}].rowid = [{fts}].rowid
where
[{fts}] match :search
order by
{order}
limit
{limit}
""".format(
dbtable=dbtable,
original=original,
columns=columns,
rank=rank,
fts=fts_table,
order=order if order else "{} desc".format(rank),
limit=20,
)
).strip()
if show_sql: if show_sql:
click.echo(sql) click.echo(sql)
return return
@ -1072,7 +1042,7 @@ def search(
table=table, table=table,
fmt=fmt, fmt=fmt,
json_cols=json_cols, json_cols=json_cols,
param=[("search", q)], param=[("query", q)],
load_extension=load_extension, load_extension=load_extension,
) )

View file

@ -1313,6 +1313,49 @@ class Table(Queryable):
) )
return self return self
def search_sql(self, columns=None, order=None, limit=None):
# Pick names for table and rank column that don't clash
original = "original_" if self.name == "original" else "original"
rank = "rank"
while rank in self.columns_dict:
rank = rank + "_"
columns_sql = "*"
if columns:
columns_sql = ", ".join("[{}]".format(c) for c in columns)
fts_table = self.detect_fts()
assert fts_table, "Full-text search is not configured for table '{}'".format(
self.name
)
return textwrap.dedent(
"""
with {original} as (
select
rowid,
{columns}
from [{dbtable}]
)
select
{original}.*,
[{fts}].rank as {rank}
from
[{original}]
join [{fts}] on [{original}].rowid = [{fts}].rowid
where
[{fts}] match :query
order by
{order}
{limit}
""".format(
dbtable=self.name,
original=original,
columns=columns_sql,
rank=rank,
fts=fts_table,
order=order or "{} desc".format(rank),
limit="limit {}".format(limit) if limit else "",
)
).strip()
def search(self, q): def search(self, q):
sql = ( sql = (
textwrap.dedent( textwrap.dedent(

View file

@ -276,3 +276,64 @@ def test_enable_fts_replace_does_nothing_if_args_the_same():
db["books"].enable_fts(["title", "author"], create_triggers=True, replace=True) db["books"].enable_fts(["title", "author"], create_triggers=True, replace=True)
# The only SQL that executed should be select statements # The only SQL that executed should be select statements
assert all(q[0].startswith("select ") for q in queries) assert all(q[0].startswith("select ") for q in queries)
@pytest.mark.parametrize(
"kwargs,expected",
[
(
{},
(
"with original as (\n"
" select\n"
" rowid,\n"
" *\n"
" from [books]\n"
")\n"
"select\n"
" original.*,\n"
" [books_fts].rank as rank\n"
"from\n"
" [original]\n"
" join [books_fts] on [original].rowid = [books_fts].rowid\n"
"where\n"
" [books_fts] match :query\n"
"order by\n"
" rank desc"
),
),
(
{"columns": ["title"], "order": "rowid", "limit": 10},
(
"with original as (\n"
" select\n"
" rowid,\n"
" [title]\n"
" from [books]\n"
")\n"
"select\n"
" original.*,\n"
" [books_fts].rank as rank\n"
"from\n"
" [original]\n"
" join [books_fts] on [original].rowid = [books_fts].rowid\n"
"where\n"
" [books_fts] match :query\n"
"order by\n"
" rowid\n"
"limit 10"
),
),
],
)
def test_search_sql(kwargs, expected):
db = Database(memory=True)
db["books"].insert(
{
"title": "Habits of Australian Marsupials",
"author": "Marlee Hawkins",
}
)
db["books"].enable_fts(["title", "author"])
sql = db["books"].search_sql(**kwargs)
assert sql == expected