Add suggestion when setting the search_path (dbcli#1513)

To be suggested autocompletion when setting the search_path, the
keyword's suggestion "search_path TO" will be suggested on a "SET"
token, only if the latter is the first token of the SQL statement (to
differentiate it from the "UPDATE foo SET"). Then, available schemas are
suggested (after the "TO" token).

Resolves dbcli#1511

Author: fbdb

AUTHORS

@@ -142,6 +142,7 @@ Contributors:
     * Fabio (3ximus)
     * Doug Harris (dougharris)
     * Jay Knight (jay-knight)
+    * fbdb
 
 Creator:
 --------

changelog.rst

@@ -7,6 +7,7 @@ Features:
     * Command line option `--init-command`
     * Provide `init-command` in the config file
     * Support dsn specific init-command in the config file
+* Add suggestion when setting the search_path
 
 Internal:
 ---------

pgcli/packages/pgliterals/pgliterals.json

@@ -227,7 +227,7 @@
         "ROWS": [],
         "SELECT": [],
         "SESSION": [],
-        "SET": [],
+        "SET": ["SEARCH_PATH TO"],
         "SHARE": [],
         "SHOW": [],
         "SIZE": [],

pgcli/packages/sqlcompletion.py

@@ -372,7 +372,15 @@ def suggest_based_on_last_token(token, stmt):
         # We're probably in a function argument list
         return _suggest_expression(token_v, stmt)
     elif token_v == "set":
+        # "set" for changing a run-time parameter
+        p = sqlparse.parse(stmt.text_before_cursor)[0]
+        is_first_token = p.token_first().value.upper() == token_v.upper()
+        if is_first_token:
+            return (Keyword(token_v.upper()),)
+
+        # E.g. 'UPDATE foo SET'
         return (Column(table_refs=stmt.get_tables(), local_tables=stmt.local_tables),)
+
     elif token_v in ("select", "where", "having", "order by", "distinct"):
         return _suggest_expression(token_v, stmt)
     elif token_v == "as":
@@ -494,6 +502,9 @@ def suggest_based_on_last_token(token, stmt):
         return tuple(suggestions)
     elif token_v in {"alter", "create", "drop"}:
         return (Keyword(token_v.upper()),)
+    elif token_v == "to":
+        # E.g. 'SET search_path TO'
+        return (Schema(),)
     elif token.is_keyword:
         # token is a keyword we haven't implemented any special handling for
         # go backwards in the query until we find one we do recognize

tests/test_sqlcompletion.py

@@ -918,3 +918,13 @@ def test_handle_unrecognized_kw_generously():
 @pytest.mark.parametrize("sql", ["ALTER ", "ALTER TABLE foo ALTER "])
 def test_keyword_after_alter(sql):
     assert Keyword("ALTER") in set(suggest_type(sql, sql))
+
+
+def test_suggestion_when_setting_search_path():
+    sql_set = "SET "
+    suggestion_set = suggest_type(sql_set, sql_set)
+    assert set(suggestion_set) == {Keyword("SET")}
+
+    sql_set_search_path_to = "SET search_path TO "
+    suggestion_set_search_path_to = suggest_type(sql_set_search_path_to, sql_set_search_path_to)
+    assert set(suggestion_set_search_path_to) == {Schema()}