where= and where_args= parameters to search() and search_sql()

Closes #441

docs/python-api.rst

@@ -2108,17 +2108,25 @@ 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 <python_api_quote_fts>` 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
 
     for article in db["articles"].search(
         "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
 

sqlite_utils/db.py

@@ -2254,6 +2254,7 @@ def search_sql(
         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 @@ def search_sql(
         :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 @@ def search_sql(
             select
                 rowid,
                 {columns}
-            from [{dbtable}]
+            from [{dbtable}]{where_clause}
         )
         select
             {columns_with_prefix}
@@ -2311,6 +2313,7 @@ def search_sql(
             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 @@ def search(
         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 @@ def search(
         :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:

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):