Repository: dbcli/pgcli

Issue:
Autocompletion when setting the search_path
As discussed [here](https://github.com/dbcli/pgcli/discussions/1509), there is no autocompletion when setting the search_path (`SET search_path TO bar`).

[As suggested](https://github.com/dbcli/pgcli/discussions/1509#discussioncomment-13276295), the first autocompletion can easily be implemented by adding `search_path TO` here: https://github.com/dbcli/pgcli/blob/435cf7f2843c1a161ad63d72b75b91d1ccba8200/pgcli/packages/pgliterals/pgliterals.json#L230


Discussion:

Patch:
diff --git a/AUTHORS b/AUTHORS
index d62e13d21..ef4a014c7 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -142,6 +142,7 @@ Contributors:
     * Fabio (3ximus)
     * Doug Harris (dougharris)
     * Jay Knight (jay-knight)
+    * fbdb
 
 Creator:
 --------
diff --git a/changelog.rst b/changelog.rst
index 10d69d1a8..c82dce084 100644
--- a/changelog.rst
+++ b/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:
 ---------
diff --git a/pgcli/packages/pgliterals/pgliterals.json b/pgcli/packages/pgliterals/pgliterals.json
index df00817a5..6828a4632 100644
--- a/pgcli/packages/pgliterals/pgliterals.json
+++ b/pgcli/packages/pgliterals/pgliterals.json
@@ -227,7 +227,7 @@
         "ROWS": [],
         "SELECT": [],
         "SESSION": [],
-        "SET": [],
+        "SET": ["SEARCH_PATH TO"],
         "SHARE": [],
         "SHOW": [],
         "SIZE": [],
diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index 9eb3ca858..9db4e31ed 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/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


Test patch:
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index 3f99d6a32..028170d58 100644
--- a/tests/test_sqlcompletion.py
+++ b/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()}