Skip to content

Commit

Permalink
enable_fts(), populate_fts() and search() methods
Browse files Browse the repository at this point in the history
  • Loading branch information
Simon Willison committed Jul 31, 2018
1 parent c446e22 commit f4907f6
Show file tree
Hide file tree
Showing 4 changed files with 100 additions and 1 deletion.
28 changes: 28 additions & 0 deletions docs/table.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import io
import os

VERSION = "0.3.1"
VERSION = "0.4"


def get_long_description():
Expand Down
34 changes: 34 additions & 0 deletions sqlite_utils/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,28 @@ def add_foreign_key(self, column, column_type, other_table, other_column):
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:
Expand All @@ -155,6 +177,18 @@ def detect_column_types(self, records):
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
Expand Down
37 changes: 37 additions & 0 deletions tests/test_enable_fts.py
Original file line number Diff line number Diff line change
@@ -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")

0 comments on commit f4907f6

Please sign in to comment.