From 8ae77a6961fed94ef2c9cc81fcfc7c81d222d9a2 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 18 Aug 2021 12:55:53 -0700
Subject: [PATCH] table.search(quote=True) parameter, refs #296

---
 docs/python-api.rst | 14 ++++++++++++--
 sqlite_utils/db.py  | 16 ++++++++++++----
 tests/test_fts.py   | 37 +++++++++++++++++++++++++++++++------
 3 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/docs/python-api.rst b/docs/python-api.rst
index 0a0e1df..20d080d 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -1843,8 +1843,15 @@ The ``.has_counts_triggers`` property shows if a table has been configured with
 
 .. _python_api_fts:
 
-Enabling full-text search
-=========================
+Full-text search
+================
+
+SQLite includes bundled extensions that implement `powerful full-text search <https://www.sqlite.org/fts5.html>`__.
+
+.. _python_api_fts_enable:
+
+Enabling full-text search for a table
+-------------------------------------
 
 You can enable full-text search on a table using ``.enable_fts(columns)``:
 
@@ -1947,6 +1954,9 @@ The ``.search()`` method also accepts the following optional parameters:
 ``offset`` integer
     Offset to use along side the limit parameter.
 
+``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:
 
 .. code-block:: python
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 306ff18..e4c1d5c 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -1924,7 +1924,13 @@ class Table(Queryable):
             )
         return self
 
-    def search_sql(self, columns=None, order_by=None, limit=None, offset=None) -> str:
+    def search_sql(
+        self,
+        columns: Optional[Iterable[str]] = None,
+        order_by: Optional[str] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+    ) -> str:
         "Return SQL string that can be used to execute searches against this table."
         # Pick names for table and rank column that don't clash
         original = "original_" if self.name == "original" else "original"
@@ -1986,19 +1992,21 @@ class Table(Queryable):
         self,
         q: str,
         order_by: Optional[str] = None,
-        columns: Optional[List[str]] = None,
+        columns: Optional[Iterable[str]] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
+        quote: bool = False,
     ) -> Generator[dict, None, None]:
         """
         Execute a search against this table using SQLite full-text search, returning a sequence of
         dictionaries for each row.
 
-        - ``q`` - words to search for
+        - ``q`` - terms to search for
         - ``order_by`` - defaults to order by rank, or specify a column here.
         - ``columns`` - list of columns to return, defaults to all columns.
         - ``limit`` - optional integer limit for returned rows.
         - ``offset`` - optional integer SQL offset.
+        - ``quote`` - apply quoting to disable any special characters in the search query
 
         See :ref:`python_api_fts_search`.
         """
@@ -2009,7 +2017,7 @@ class Table(Queryable):
                 limit=limit,
                 offset=offset,
             ),
-            {"query": q},
+            {"query": self.db.quote_fts(q) if quote else q},
         )
         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 7a414da..6aa6fbe 100644
--- a/tests/test_fts.py
+++ b/tests/test_fts.py
@@ -503,13 +503,38 @@ def test_search_sql(kwargs, fts, expected):
     assert sql == expected
 
 
-def test_quote_fts_query(fresh_db):
-
+@pytest.mark.parametrize(
+    "input,expected",
+    (
+        ("dog", '"dog"'),
+        ("cat,", '"cat,"'),
+        ("cat's", '"cat\'s"'),
+        ("dog.", '"dog."'),
+        ("cat dog", '"cat" "dog"'),
+        # If a phrase is already double quoted, leave it so
+        ('"cat dog"', '"cat dog"'),
+        ('"cat dog" fish', '"cat dog" "fish"'),
+        # Sensibly handle unbalanced double quotes
+        ('cat"', '"cat"'),
+        ('"cat dog" "fish', '"cat dog" "fish"'),
+    ),
+)
+def test_quote_fts_query(fresh_db, input, expected):
     table = fresh_db["searchable"]
     table.insert_all(search_records)
     table.enable_fts(["text", "country"])
-
-    query = "cat's"
-    result = fresh_db.quote_fts(query)
+    quoted = fresh_db.quote_fts(input)
+    assert quoted == expected
     # Executing query does not crash.
-    list(table.search(result))
+    list(table.search(quoted))
+
+
+def test_search_quote(fresh_db):
+    table = fresh_db["searchable"]
+    table.insert_all(search_records)
+    table.enable_fts(["text", "country"])
+    query = "cat's"
+    with pytest.raises(sqlite3.OperationalError):
+        list(table.search(query))
+    # No exception with quote=True
+    list(table.search(query, quote=True))