From 1b09538bc6c1fda773590f3e600993ef06591041 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Tue, 14 Jun 2022 14:54:35 -0700 Subject: [PATCH] where= and where_args= parameters to search() and search_sql() Closes #441 --- docs/python-api.rst | 12 +++++-- sqlite_utils/db.py | 20 ++++++++++-- tests/test_fts.py | 76 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 104 insertions(+), 4 deletions(-) diff --git a/docs/python-api.rst b/docs/python-api.rst index ca98317..1a8c989 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -2108,10 +2108,16 @@ The ``.search()`` method also accepts the following optional parameters: ``offset`` integer Offset to use along side the limit parameter. +``where`` string + Extra SQL fragment for the WHERE clause + +``where_args`` dictionary + Arguments to use for ``:param`` placeholders in the extra WHERE clause + ``quote`` bool Apply :ref:`FTS quoting rules ` to the search query, disabling advanced query syntax in a way that avoids surprising errors. -To return just the title and published columns for three matches for ``"dog"`` ordered by ``published`` with the most recent first, use the following: +To return just the title and published columns for three matches for ``"dog"`` where the ``id`` is greater than 10 ordered by ``published`` with the most recent first, use the following: .. code-block:: python @@ -2119,6 +2125,8 @@ To return just the title and published columns for three matches for ``"dog"`` o "dog", order_by="published desc", limit=3, + where="id > :min_id", + where_args={"min_id": 10}, columns=["title", "published"] ): print(article) @@ -2128,7 +2136,7 @@ To return just the title and published columns for three matches for ``"dog"`` o Building SQL queries with table.search_sql() -------------------------------------------- -You can generate the SQL query that would be used for a search using the ``table.search_sql()`` method. It takes the same arguments as ``table.search()`` with the exception of the search query itself, since the returned SQL includes a parameter that can be used for the search. +You can generate the SQL query that would be used for a search using the ``table.search_sql()`` method. It takes the same arguments as ``table.search()``, with the exception of the search query and the ``where_args`` parameter, since those should be provided when the returned SQL is executed. .. code-block:: python diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index cbf59d6..7a06304 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -2254,6 +2254,7 @@ class Table(Queryable): order_by: Optional[str] = None, limit: Optional[int] = None, offset: Optional[int] = None, + where: Optional[str] = None, ) -> str: """ " Return SQL string that can be used to execute searches against this table. @@ -2262,6 +2263,7 @@ class Table(Queryable): :param order_by: Column or SQL expression to sort by :param limit: SQL limit :param offset: SQL offset + :param where: Extra SQL fragment for the WHERE clause """ # Pick names for table and rank column that don't clash original = "original_" if self.name == "original" else "original" @@ -2283,7 +2285,7 @@ class Table(Queryable): select rowid, {columns} - from [{dbtable}] + from [{dbtable}]{where_clause} ) select {columns_with_prefix} @@ -2311,6 +2313,7 @@ class Table(Queryable): limit_offset += " offset {}".format(offset) return sql.format( dbtable=self.name, + where_clause="\n where {}".format(where) if where else "", original=original, columns=columns_sql, columns_with_prefix=columns_with_prefix_sql, @@ -2326,6 +2329,8 @@ class Table(Queryable): columns: Optional[Iterable[str]] = None, limit: Optional[int] = None, offset: Optional[int] = None, + where: str = None, + where_args: Optional[Union[Iterable, dict]] = None, quote: bool = False, ) -> Generator[dict, None, None]: """ @@ -2337,18 +2342,29 @@ class Table(Queryable): :param columns: List of columns to return, defaults to all columns. :param limit: Optional integer limit for returned rows. :param offset: Optional integer SQL offset. + :param where: Extra SQL fragment for the WHERE clause + :param where_args: Arguments to use for :param placeholders in the extra WHERE clause :param quote: Apply quoting to disable any special characters in the search query See :ref:`python_api_fts_search`. """ + args = {"query": self.db.quote_fts(q) if quote else q} + if where_args and "query" in where_args: + raise ValueError( + "'query' is a reserved key and cannot be passed to where_args for .search()" + ) + if where_args: + args.update(where_args) + cursor = self.db.execute( self.search_sql( order_by=order_by, columns=columns, limit=limit, offset=offset, + where=where, ), - {"query": self.db.quote_fts(q) if quote else q}, + args, ) columns = [c[0] for c in cursor.description] for row in cursor: diff --git a/tests/test_fts.py b/tests/test_fts.py index 24980f7..53a7bc2 100644 --- a/tests/test_fts.py +++ b/tests/test_fts.py @@ -94,6 +94,38 @@ def test_search_limit_offset(fresh_db): ) +@pytest.mark.parametrize("fts_version", ("FTS4", "FTS5")) +def test_search_where(fresh_db, fts_version): + table = fresh_db["t"] + table.insert_all(search_records) + table.enable_fts(["text", "country"], fts_version=fts_version) + results = list( + table.search("are", where="country = :country", where_args={"country": "Japan"}) + ) + assert results == [ + { + "rowid": 1, + "text": "tanuki are running tricksters", + "country": "Japan", + "not_searchable": "foo", + } + ] + + +def test_search_where_args_disallows_query(fresh_db): + table = fresh_db["t"] + with pytest.raises(ValueError) as ex: + list( + table.search( + "x", where="author = :query", where_args={"query": "not allowed"} + ) + ) + assert ( + ex.value.args[0] + == "'query' is a reserved key and cannot be passed to where_args for .search()" + ) + + def test_enable_fts_table_names_containing_spaces(fresh_db): table = fresh_db["test"] table.insert({"column with spaces": "in its name"}) @@ -415,6 +447,28 @@ def test_enable_fts_error_message_on_views(): "limit 10" ), ), + ( + {"where": "author = :author"}, + "FTS5", + ( + "with original as (\n" + " select\n" + " rowid,\n" + " *\n" + " from [books]\n" + " where author = :author\n" + ")\n" + "select\n" + " [original].*\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" + " [books_fts].rank" + ), + ), ( {"columns": ["title"]}, "FTS4", @@ -480,6 +534,28 @@ def test_enable_fts_error_message_on_views(): "limit 2" ), ), + ( + {"where": "author = :author"}, + "FTS4", + ( + "with original as (\n" + " select\n" + " rowid,\n" + " *\n" + " from [books]\n" + " where author = :author\n" + ")\n" + "select\n" + " [original].*\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_bm25(matchinfo([books_fts], 'pcnalx'))" + ), + ), ], ) def test_search_sql(kwargs, fts, expected):