diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d91317..201e7f4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ### Bug Fixes - Expand `~` in configured log file paths before opening the log. +- Avoid completion refresh crashes when a relation references a function that is unavailable in the current session. ### Internal diff --git a/litecli/completion_refresher.py b/litecli/completion_refresher.py index 4e76faa..4052fe2 100644 --- a/litecli/completion_refresher.py +++ b/litecli/completion_refresher.py @@ -125,8 +125,11 @@ def refresh_schemata(completer: SQLCompleter, executor: SQLExecute) -> None: @refresher("tables") def refresh_tables(completer: SQLCompleter, executor: SQLExecute) -> None: + # Register the names of all relations first so that relations whose columns + # cannot be introspected in this session (e.g. views that reference a + # function provided by a loadable extension) still complete by name. + completer.extend_relations(list(executor.tables()), kind="tables") table_cols = list(executor.table_columns()) - completer.extend_relations(table_cols, kind="tables") completer.extend_columns(table_cols, kind="tables") diff --git a/litecli/sqlexecute.py b/litecli/sqlexecute.py index 5e52b73..981ecbf 100644 --- a/litecli/sqlexecute.py +++ b/litecli/sqlexecute.py @@ -49,6 +49,12 @@ class SQLExecute(object): ORDER BY tableName, columnName """ + relation_columns_query = """ + SELECT p.name as columnName + FROM pragma_table_info(?) p + ORDER BY columnName + """ + indexes_query = """ SELECT name, sql FROM sqlite_master @@ -194,9 +200,27 @@ def table_columns(self) -> Generator[tuple[str, str], None, None]: return with closing(self.conn.cursor()) as cur: _logger.debug("Columns Query. sql: %r", self.table_columns_query) - cur.execute(self.table_columns_query) - for row in cur: - yield row + try: + rows = list(cur.execute(self.table_columns_query)) + except OperationalError: + # A relation in the database (e.g. a view) can reference a + # function that is not available in this session, such as one + # provided by a loadable extension. Introspect each relation + # individually and skip the ones that cannot be introspected. + _logger.debug("Columns Query failed. Introspecting relations individually.") + cur.execute(self.tables_query) + relations = [row[0] for row in cur] + for name in relations: + with closing(self.conn.cursor()) as columns_cur: + try: + columns_cur.execute(self.relation_columns_query, (name,)) + except OperationalError: + _logger.debug("Unable to introspect the columns of %s.", name) + continue + for (column,) in columns_cur: + yield (name, column) + return + yield from rows def databases(self) -> Generator[str, None, None]: if not self.conn: diff --git a/tests/test_completion_refresher.py b/tests/test_completion_refresher.py index 32da6bd..fa63bf4 100644 --- a/tests/test_completion_refresher.py +++ b/tests/test_completion_refresher.py @@ -93,3 +93,37 @@ def test_refresh_with_callbacks(refresher): refresher.refresh(sqlexecute, callbacks) time.sleep(1) # Wait for the thread to work. assert callbacks[0].call_count == 1 + + +def test_refresh_with_view_referencing_unavailable_function(refresher, tmp_path): + """A view that references a function provided by a loadable extension + must not crash the background completion refresh. All relation names + and the columns of the introspectable relations must be completed. + + :param refresher: + + """ + from litecli.sqlexecute import SQLExecute, sqlite3 + + db_path = str(tmp_path / "completion_refresh.db") + conn = sqlite3.connect(db_path) + conn.execute("create table products(product_id integer primary key)") + conn.create_function("unavailable_ext_func", 1, lambda x: b"") + conn.execute("create view v_ext_hashed as select hex(unavailable_ext_func(product_id)) from products") + conn.commit() + conn.close() + + executor = SQLExecute(db_path) + refreshed = [] + refresher.refresh(executor, refreshed.append, {}) + + deadline = time.time() + 5 + while refresher.is_refreshing() and time.time() < deadline: + time.sleep(0.05) + time.sleep(0.1) + + assert refreshed, "the completion refresh did not complete" + completer = refreshed[0] + assert "products" in completer.all_completions + assert "v_ext_hashed" in completer.all_completions + assert "product_id" in completer.all_completions diff --git a/tests/test_sqlexecute.py b/tests/test_sqlexecute.py index 3fdaef1..bf87b40 100644 --- a/tests/test_sqlexecute.py +++ b/tests/test_sqlexecute.py @@ -1,6 +1,7 @@ # coding=UTF-8 import os +from contextlib import closing from typing import Any import pytest @@ -63,6 +64,25 @@ def test_table_and_columns_query(executor): assert set(executor.table_columns()) == set([("a", "x"), ("a", "y"), ("b", "z"), ("t", "t")]) +@dbtest +def test_table_columns_with_function_only_available_in_another_session(executor, connection): + run(executor, "create table products(product_id integer primary key, product_name text)") + + # The view references a function that is only registered in the session of + # the `connection` fixture, as a function provided by a loadable extension + # would be. It cannot be introspected from the executor's session. + with closing(connection.cursor()) as cur: + connection.create_function("unavailable_ext_func", 1, lambda x: b"") + cur.execute("create view v_ext_hashed as select hex(unavailable_ext_func(product_id)) from products") + + assert ("v_ext_hashed",) in set(executor.tables()) + + columns = set(executor.table_columns()) + assert ("products", "product_id") in columns + assert ("products", "product_name") in columns + assert not any(table == "v_ext_hashed" for table, _ in columns) + + @dbtest def test_database_list(executor): databases = executor.databases()