From f4907f6df58d822dfb67660b982a9081b39a06fb Mon Sep 17 00:00:00 2001
From: Simon Willison <simonw@eventbrite.com>
Date: Tue, 31 Jul 2018 09:19:05 -0700
Subject: [PATCH] enable_fts(), populate_fts() and search() methods

---
 docs/table.rst           | 28 ++++++++++++++++++++++++++++
 setup.py                 |  2 +-
 sqlite_utils/db.py       | 34 ++++++++++++++++++++++++++++++++++
 tests/test_enable_fts.py | 37 +++++++++++++++++++++++++++++++++++++
 4 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_enable_fts.py

diff --git a/docs/table.rst b/docs/table.rst
index 27df299..0ae25d3 100644
--- a/docs/table.rst
+++ b/docs/table.rst
@@ -116,3 +116,31 @@ The ``.schema`` property outputs the table's schema as a SQL string::
         FOREIGN KEY ("qSiteInfo") REFERENCES [qSiteInfo](id),
         FOREIGN KEY ("qCareAssistant") REFERENCES [qCareAssistant](id),
         FOREIGN KEY ("qLegalStatus") REFERENCES [qLegalStatus](id))
+
+Enabling full-text search
+=========================
+
+You can enable full-text search on a table using ``.enable_fts(columns)``:
+
+.. code-block:: python
+
+    dogs.enable_fts(["name", "twitter"])
+
+You can then run searches using the ``.search()`` method:
+
+.. code-block:: python
+
+    rows = dogs.search("cleo")
+
+If you insert additioal records into the table you will need to refresh the search index using ``populate_fts()``:
+
+.. code-block:: python
+
+    dogs.insert({
+        "id": 2,
+        "name": "Marnie",
+        "twitter": "MarnieTheDog",
+        "age": 16,
+        "is_good_dog": True,
+    }, pk="id")
+    dogs.populate_fts(["name", "twitter"])
diff --git a/setup.py b/setup.py
index fffd489..a044e5f 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 import io
 import os
 
-VERSION = "0.3.1"
+VERSION = "0.4"
 
 
 def get_long_description():
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 0ceb7a8..b5fbf26 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -133,6 +133,28 @@ class Table:
         self.db.conn.commit()
         return result
 
+    def enable_fts(self, columns):
+        "Enables FTS on the specified columns"
+        sql = """
+            CREATE VIRTUAL TABLE "{table}_fts" USING FTS4 (
+                {columns},
+                content="{table}"
+            );
+        """.format(
+            table=self.name, columns=", ".join(columns)
+        )
+        self.db.conn.executescript(sql)
+        self.populate_fts(columns)
+
+    def populate_fts(self, columns):
+        sql = """
+            INSERT INTO "{table}_fts" (rowid, {columns})
+                SELECT rowid, {columns} FROM {table};
+        """.format(
+            table=self.name, columns=", ".join(columns)
+        )
+        self.db.conn.executescript(sql)
+
     def detect_column_types(self, records):
         all_column_types = {}
         for record in records:
@@ -155,6 +177,18 @@ class Table:
             column_types[key] = t
         return column_types
 
+    def search(self, q):
+        sql = """
+            select * from {table} where rowid in (
+                select rowid from [{table}_fts]
+                where [{table}_fts] match :search
+            )
+            order by rowid
+        """.format(
+            table=self.name
+        )
+        return self.db.conn.execute(sql, (q,)).fetchall()
+
     def insert(self, record, pk=None, foreign_keys=None, upsert=False):
         return self.insert_all(
             [record], pk=pk, foreign_keys=foreign_keys, upsert=upsert
diff --git a/tests/test_enable_fts.py b/tests/test_enable_fts.py
new file mode 100644
index 0000000..8a95042
--- /dev/null
+++ b/tests/test_enable_fts.py
@@ -0,0 +1,37 @@
+from .fixtures import fresh_db
+import pytest
+
+search_records = [
+    {"text": "tanuki are tricksters", "country": "Japan", "not_searchable": "foo"},
+    {"text": "racoons are trash pandas", "country": "USA", "not_searchable": "bar"},
+]
+
+
+def test_enable_fts(fresh_db):
+    table = fresh_db["searchable"]
+    table.insert_all(search_records)
+    assert ["searchable"] == fresh_db.tables
+    table.enable_fts(["text", "country"])
+    assert [
+        "searchable",
+        "searchable_fts",
+        "searchable_fts_segments",
+        "searchable_fts_segdir",
+        "searchable_fts_docsize",
+        "searchable_fts_stat",
+    ] == fresh_db.tables
+    assert [("tanuki are tricksters", "Japan", "foo")] == table.search("tanuki")
+    assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa")
+    assert [] == table.search("bar")
+
+
+def test_populate_fts(fresh_db):
+    table = fresh_db["populatable"]
+    table.insert(search_records[0])
+    table.enable_fts(["text", "country"])
+    assert [] == table.search("trash pandas")
+    table.insert(search_records[1])
+    assert [] == table.search("trash pandas")
+    # Now run populate_fts to make this record available
+    table.populate_fts(["text", "country"])
+    assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa")