From de39e8db1ee8b18755b9d83c69371a65664106fd Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Tue, 3 Nov 2020 14:46:18 -0800 Subject: [PATCH] Refactored to table.search_sql() method, added --limit --- docs/cli.rst | 6 ++--- sqlite_utils/cli.py | 54 +++++++++------------------------------ sqlite_utils/db.py | 43 ++++++++++++++++++++++++++++++++ tests/test_fts.py | 61 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 119 insertions(+), 45 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index 5d20116..2d7c638 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -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 +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:: $ 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] join [documents_fts] on [original].rowid = [documents_fts].rowid where - [documents_fts] match :search + [documents_fts] match :query order by rank desc - limit - 20 .. _cli_vacuum: diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 15fc4b2..15c78f6 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -7,7 +7,6 @@ import hashlib import pathlib import sqlite_utils from sqlite_utils.db import AlterError -import textwrap import itertools import json import os @@ -979,6 +978,12 @@ def query( @click.argument("q") @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( + "--limit", + type=int, + default=20, + help="Number of rows to return, default 20, set to 0 for unlimited", +) @click.option( "--sql", "show_sql", is_flag=True, help="Show SQL query that would be run" ) @@ -993,6 +998,7 @@ def search( order, show_sql, column, + limit, nl, arrays, csv, @@ -1009,55 +1015,19 @@ def search( table_obj = db[dbtable] if not table_obj.exists(): raise click.ClickException("Table '{}' does not exist".format(dbtable)) - fts_table = table_obj.detect_fts() - if not fts_table: + if not table_obj.detect_fts(): raise click.ClickException( "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: # Check they all exist + table_columns = table_obj.columns_dict for c in column: - if c not in table_obj.columns_dict: + if c not in table_columns: raise click.ClickException( "Table '{}' has no column '{}".format(dbtable, c) ) - columns = ", ".join("[{}]".format(c) for c in column) - 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() + sql = table_obj.search_sql(columns=column, order=order, limit=limit) if show_sql: click.echo(sql) return @@ -1072,7 +1042,7 @@ def search( table=table, fmt=fmt, json_cols=json_cols, - param=[("search", q)], + param=[("query", q)], load_extension=load_extension, ) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 3600e88..cd6a5a6 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1313,6 +1313,49 @@ class Table(Queryable): ) 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): sql = ( textwrap.dedent( diff --git a/tests/test_fts.py b/tests/test_fts.py index 359d911..ddcf1bd 100644 --- a/tests/test_fts.py +++ b/tests/test_fts.py @@ -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) # The only SQL that executed should be select statements 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