diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 7c09578..0b51517 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -265,7 +265,13 @@ def create_index(path, table, column, name, unique, if_not_exists): @click.argument("column", nargs=-1, required=True) @click.option("--fts4", help="Use FTS4", default=False, is_flag=True) @click.option("--fts5", help="Use FTS5", default=False, is_flag=True) -def enable_fts(path, table, column, fts4, fts5): +@click.option( + "--create-triggers", + help="Create triggers to update the FTS tables when the parent table changes.", + default=False, + is_flag=True, +) +def enable_fts(path, table, column, fts4, fts5, create_triggers): "Enable FTS for specific table and columns" fts_version = "FTS5" if fts4 and fts5: @@ -275,7 +281,9 @@ def enable_fts(path, table, column, fts4, fts5): fts_version = "FTS4" db = sqlite_utils.Database(path) - db[table].enable_fts(column, fts_version=fts_version) + db[table].enable_fts( + column, fts_version=fts_version, create_triggers=create_triggers + ) @cli.command(name="populate-fts") diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 6247428..b7837c7 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -646,7 +646,7 @@ class Table(Queryable): return self def drop(self): - self.db.conn.execute("DROP TABLE {}".format(self.name)) + self.db.conn.execute("DROP TABLE [{}]".format(self.name)) def guess_foreign_table(self, column): column = column.lower() @@ -710,12 +710,12 @@ class Table(Queryable): ) self.db.add_foreign_keys([(self.name, column, other_table, other_column)]) - def enable_fts(self, columns, fts_version="FTS5"): - "Enables FTS on the specified columns" + def enable_fts(self, columns, fts_version="FTS5", create_triggers=False): + "Enables FTS on the specified columns." sql = """ - CREATE VIRTUAL TABLE "{table}_fts" USING {fts_version} ( + CREATE VIRTUAL TABLE [{table}_fts] USING {fts_version} ( {columns}, - content="{table}" + content=[{table}] ); """.format( table=self.name, @@ -724,12 +724,34 @@ class Table(Queryable): ) self.db.conn.executescript(sql) self.populate_fts(columns) + + if create_triggers: + old_cols = ", ".join("old.[{}]".format(c) for c in columns) + new_cols = ", ".join("new.[{}]".format(c) for c in columns) + triggers = """ + CREATE TRIGGER [{table}_ai] AFTER INSERT ON [{table}] BEGIN + INSERT INTO [{table}_fts] (rowid, {columns}) VALUES (new.rowid, {new_cols}); + END; + CREATE TRIGGER [{table}_ad] AFTER DELETE ON [{table}] BEGIN + INSERT INTO [{table}_fts] ([{table}_fts], rowid, {columns}) VALUES('delete', old.rowid, {old_cols}); + END; + CREATE TRIGGER [{table}_au] AFTER UPDATE ON [{table}] BEGIN + INSERT INTO [{table}_fts] ([{table}_fts], rowid, {columns}) VALUES('delete', old.rowid, {old_cols}); + INSERT INTO [{table}_fts] (rowid, {columns}) VALUES (new.rowid, {new_cols}); + END; + """.format( + table=self.name, + columns=", ".join("[{}]".format(c) for c in columns), + old_cols=old_cols, + new_cols=new_cols, + ) + self.db.conn.executescript(triggers) return self def populate_fts(self, columns): sql = """ - INSERT INTO "{table}_fts" (rowid, {columns}) - SELECT rowid, {columns} FROM {table}; + INSERT INTO [{table}_fts] (rowid, {columns}) + SELECT rowid, {columns} FROM [{table}]; """.format( table=self.name, columns=", ".join(columns) ) @@ -738,21 +760,20 @@ class Table(Queryable): def detect_fts(self): "Detect if table has a corresponding FTS virtual table and return it" - rows = self.db.conn.execute( - """ + sql = """ SELECT name FROM sqlite_master WHERE rootpage = 0 AND ( - sql LIKE '%VIRTUAL TABLE%USING FTS%content="{table}"%' + sql LIKE '%VIRTUAL TABLE%USING FTS%content=%{table}%' OR ( tbl_name = "{table}" AND sql LIKE '%VIRTUAL TABLE%USING FTS%' ) ) """.format( - table=self.name - ) - ).fetchall() + table=self.name + ) + rows = self.db.conn.execute(sql).fetchall() if len(rows) == 0: return None else: @@ -796,7 +817,7 @@ class Table(Queryable): def search(self, q): sql = """ - select * from {table} where rowid in ( + select * from "{table}" where rowid in ( select rowid from [{table}_fts] where [{table}_fts] match :search ) @@ -1144,7 +1165,7 @@ class View(Queryable): ) def drop(self): - self.db.conn.execute("DROP VIEW {}".format(self.name)) + self.db.conn.execute("DROP VIEW [{}]".format(self.name)) def chunks(sequence, size): diff --git a/tests/test_cli.py b/tests/test_cli.py index ddb936b..ca91bdb 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -315,6 +315,7 @@ def test_index_foreign_keys(db_path): db = Database(db_path) assert [] == db["books"].indexes result = CliRunner().invoke(cli.cli, ["index-foreign-keys", db_path]) + assert 0 == result.exit_code assert [["author_id"], ["author_name_ref"]] == [ i.columns for i in db["books"].indexes ] @@ -328,6 +329,43 @@ def test_enable_fts(db_path): assert 0 == result.exit_code assert "Gosh_fts" == Database(db_path)["Gosh"].detect_fts() + # Table names with restricted chars are handled correctly. + # colons and dots are restricted characters for table names. + Database(db_path)["http://example.com"].create({"c1": str, "c2": str, "c3": str}) + assert None == Database(db_path)["http://example.com"].detect_fts() + result = CliRunner().invoke( + cli.cli, ["enable-fts", db_path, "http://example.com", "c1", "--fts4"] + ) + assert 0 == result.exit_code + assert ( + "http://example.com_fts" == Database(db_path)["http://example.com"].detect_fts() + ) + Database(db_path)["http://example.com"].drop() + + +def test_enable_fts_with_triggers(db_path): + Database(db_path)["Gosh"].insert_all([{"c1": "baz"}]) + exit_code = ( + CliRunner() + .invoke( + cli.cli, + ["enable-fts", db_path, "Gosh", "c1", "--fts4", "--create-triggers"], + ) + .exit_code + ) + assert 0 == exit_code + + def search(q): + return ( + Database(db_path) + .conn.execute("select c1 from Gosh_fts where c1 match ?", [q]) + .fetchall() + ) + + assert [("baz",)] == search("baz") + Database(db_path)["Gosh"].insert_all([{"c1": "martha"}]) + assert [("martha",)] == search("martha") + def test_populate_fts(db_path): Database(db_path)["Gosh"].insert_all([{"c1": "baz"}]) diff --git a/tests/test_enable_fts.py b/tests/test_enable_fts.py index 6cefe54..7f58916 100644 --- a/tests/test_enable_fts.py +++ b/tests/test_enable_fts.py @@ -22,6 +22,26 @@ def test_enable_fts(fresh_db): assert [] == table.search("bar") +def test_enable_fts_escape_table_names(fresh_db): + # Table names with restricted chars are handled correctly. + # colons and dots are restricted characters for table names. + table = fresh_db["http://example.com"] + table.insert_all(search_records) + assert ["http://example.com"] == fresh_db.table_names() + table.enable_fts(["text", "country"], fts_version="FTS4") + assert [ + "http://example.com", + "http://example.com_fts", + "http://example.com_fts_segments", + "http://example.com_fts_segdir", + "http://example.com_fts_docsize", + "http://example.com_fts_stat", + ] == fresh_db.table_names() + 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]) @@ -34,6 +54,19 @@ def test_populate_fts(fresh_db): assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa") +def test_populate_fts_escape_table_names(fresh_db): + # Restricted characters such as colon and dots should be escaped. + table = fresh_db["http://example.com"] + table.insert(search_records[0]) + table.enable_fts(["text", "country"], fts_version="FTS4") + 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") + + def test_optimize_fts(fresh_db): for fts_version in ("4", "5"): table_name = "searchable_{}".format(fts_version) @@ -48,3 +81,14 @@ def test_optimize_fts(fresh_db): "searchable_5_fts", ): fresh_db[table_name].optimize() + + +def test_enable_fts_w_triggers(fresh_db): + table = fresh_db["searchable"] + table.insert(search_records[0]) + table.enable_fts(["text", "country"], fts_version="FTS4", create_triggers=True) + assert [("tanuki are tricksters", "Japan", "foo")] == table.search("tanuki") + table.insert(search_records[1]) + # Triggers will auto-populate FTS virtual table, not need to call populate_fts() + assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa") + assert [] == table.search("bar")