Add suggestion when setting the search_path

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

@@ -276,7 +276,7 @@
     return (Keyword(), Special())
 
 
 def suggest_based_on_last_token(token, stmt):
     if isinstance(token, str):
         token_v = token.lower()
     elif isinstance(token, Comparison):
@@ -372,7 +372,15 @@
         # 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 @@
         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()}