Add suggestion when setting the search_path

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 #1511

Author: fbdb

AUTHORS

@@ -141,6 +141,7 @@ Contributors:
     * Josh Lynch (josh-lynch)
     * Fabio (3ximus)
     * Doug Harris (dougharris)
+    * 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()}