Allow defining a json file with preferred aliases (#1382)

* fix psycopg.sql.Identifier in \ev handling (#1384)

* Allow defining a json file with preferred aliases

At $WORK we have a lot of tables with names like `foo_noun_verb` or
`foo_noun_related-noun_verb` and so while the default aliasing is very
helpful for shortening unwieldy names we do end up with lots of aliases
like `LEFT JOIN fnv on fnv2.id = fnv.fnv2_id`

This change will allow defining a json file of preferred aliases

```
> cat ~/.config/pgcli/aliases.json
{
    "foo_user": "user",
    "foo_user_group": "user_group"
}
```

so the alias suggestion for `SELECT * FROM foo_user` will be `SELECT * FROM foo_user AS user`
instead of the default `SELECT * FROM foo_user AS fu`

* When cannot open or parse alias_map_file raise error

Raise a (hopefully) helpful exception when the alias_map_file cannot be
parsed or does not exist

* Add tests for load_alias_map_file

* Add tests for generate_alias

* Update AUTHORS file

* Remove comment.

Discussed this on the PR with a project maintainer

---------

Co-authored-by: Andy Schoenberger <akschoenberger@gmail.com>
Co-authored-by: Rob B <rob@example.com>
Co-authored-by: Irina Truong <i.chernyavska@gmail.com>

6 files changed, 116 insertions, 1 deletions

diff --git a/AUTHORS b/AUTHORS
index 9b24c8f6..a4eac9de 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -128,6 +128,7 @@ Contributors:
     * Andy Schoenberger (andyscho)
     * Damien Baty (dbaty)
     * blag
+    * Rob Berry (rob-b)
 
 Creator:
 --------
diff --git a/changelog.rst b/changelog.rst
index 8584bfec..13c1080d 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -28,6 +28,9 @@ Bug fixes:
 * Fix explain mode when used with `expand`, `auto_expand`, or `--explain-vertical-output` ([issue 1393](https://github.com/dbcli/pgcli/issues/1393)).
 * Fix sql-insert format emits NULL as 'None' ([issue 1408](https://github.com/dbcli/pgcli/issues/1408)).
 * Improve check for prompt-toolkit 3.0.6 ([issue 1416](https://github.com/dbcli/pgcli/issues/1416)).
+* Allow specifying an `alias_map_file` in the config that will use
+  predetermined table aliases instead of generating aliases programmatically on
+  the fly
 
 3.5.0 (2022/09/15):
 ===================
diff --git a/pgcli/main.py b/pgcli/main.py
index 8459178b..f95c8000 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -279,6 +279,7 @@ class PGCli:
             "single_connection": single_connection,
             "less_chatty": less_chatty,
             "keyword_casing": keyword_casing,
+            "alias_map_file": c["main"]["alias_map_file"] or None,
         }
 
         completer = PGCompleter(
diff --git a/pgcli/pgclirc b/pgcli/pgclirc
index 2d0563f1..51f7eae9 100644
--- a/pgcli/pgclirc
+++ b/pgcli/pgclirc
@@ -56,6 +56,14 @@ auto_retry_closed_connection = True
 # If set to True, table suggestions will include a table alias
 generate_aliases = False
 
+# Path to a json file that specifies specific table aliases to use when generate_aliases is set to True
+# the format for this file should be:
+# {
+#     "some_table_name": "desired_alias",
+#     "some_other_table_name": "another_alias"
+# }
+alias_map_file =
+
 # log_file location.
 # In Unix/Linux: ~/.config/pgcli/log
 # In Windows: %USERPROFILE%\AppData\Local\dbcli\pgcli\log
diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py
index f2000acf..17fc5405 100644
--- a/pgcli/pgcompleter.py
+++ b/pgcli/pgcompleter.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import re
 from itertools import count, repeat, chain
@@ -61,18 +62,38 @@ arg_default_type_strip_regex = re.compile(r"::[\w\.]+(\[\])?$")
 normalize_ref = lambda ref: ref if ref[0] == '"' else '"' + ref.lower() + '"'
 
 
-def generate_alias(tbl):
+def generate_alias(tbl, alias_map=None):
     """Generate a table alias, consisting of all upper-case letters in
     the table name, or, if there are no upper-case letters, the first letter +
     all letters preceded by _
     param tbl - unescaped name of the table to alias
     """
+    if alias_map and tbl in alias_map:
+        return alias_map[tbl]
     return "".join(
         [l for l in tbl if l.isupper()]
         or [l for l, prev in zip(tbl, "_" + tbl) if prev == "_" and l != "_"]
     )
 
 
+class InvalidMapFile(ValueError):
+    pass
+
+
+def load_alias_map_file(path):
+    try:
+        with open(path) as fo:
+            alias_map = json.load(fo)
+    except FileNotFoundError as err:
+        raise InvalidMapFile(
+            f"Cannot read alias_map_file - {err.filename} does not exist"
+        )
+    except json.JSONDecodeError:
+        raise InvalidMapFile(f"Cannot read alias_map_file - {path} is not valid json")
+    else:
+        return alias_map
+
+
 class PGCompleter(Completer):
     # keywords_tree: A dict mapping keywords to well known following keywords.
     # e.g. 'CREATE': ['TABLE', 'USER', ...],
@@ -100,6 +121,11 @@ class PGCompleter(Completer):
         self.call_arg_oneliner_max = settings.get("call_arg_oneliner_max", 2)
         self.search_path_filter = settings.get("search_path_filter")
         self.generate_aliases = settings.get("generate_aliases")
+        alias_map_file = settings.get("alias_map_file")
+        if alias_map_file is not None:
+            self.alias_map = load_alias_map_file(alias_map_file)
+        else:
+            self.alias_map = None
         self.casing_file = settings.get("casing_file")
         self.insert_col_skip_patterns = [
             re.compile(pattern)
diff --git a/tests/test_pgcompleter.py b/tests/test_pgcompleter.py
new file mode 100644
index 00000000..909fa0b7
--- /dev/null
+++ b/tests/test_pgcompleter.py
@@ -0,0 +1,76 @@
+import pytest
+from pgcli import pgcompleter
+
+
+def test_load_alias_map_file_missing_file():
+    with pytest.raises(
+        pgcompleter.InvalidMapFile,
+        match=r"Cannot read alias_map_file - /path/to/non-existent/file.json does not exist$",
+    ):
+        pgcompleter.load_alias_map_file("/path/to/non-existent/file.json")
+
+
+def test_load_alias_map_file_invalid_json(tmp_path):
+    fpath = tmp_path / "foo.json"
+    fpath.write_text("this is not valid json")
+    with pytest.raises(pgcompleter.InvalidMapFile, match=r".*is not valid json$"):
+        pgcompleter.load_alias_map_file(str(fpath))
+
+
+@pytest.mark.parametrize(
+    "table_name, alias",
+    [
+        ("SomE_Table", "SET"),
+        ("SOmeTabLe", "SOTL"),
+        ("someTable", "T"),
+    ],
+)
+def test_generate_alias_uses_upper_case_letters_from_name(table_name, alias):
+    assert pgcompleter.generate_alias(table_name) == alias
+
+
+@pytest.mark.parametrize(
+    "table_name, alias",
+    [
+        ("some_tab_le", "stl"),
+        ("s_ome_table", "sot"),
+        ("sometable", "s"),
+    ],
+)
+def test_generate_alias_uses_first_char_and_every_preceded_by_underscore(
+    table_name, alias
+):
+    assert pgcompleter.generate_alias(table_name) == alias
+
+
+@pytest.mark.parametrize(
+    "table_name, alias_map, alias",
+    [
+        ("some_table", {"some_table": "my_alias"}, "my_alias"),
+    ],
+)
+def test_generate_alias_can_use_alias_map(table_name, alias_map, alias):
+    assert pgcompleter.generate_alias(table_name, alias_map) == alias
+
+
+@pytest.mark.parametrize(
+    "table_name, alias_map, alias",
+    [
+        ("SomeTable", {"SomeTable": "my_alias"}, "my_alias"),
+    ],
+)
+def test_generate_alias_prefers_alias_over_upper_case_name(
+    table_name, alias_map, alias
+):
+    assert pgcompleter.generate_alias(table_name, alias_map) == alias
+
+
+@pytest.mark.parametrize(
+    "table_name, alias",
+    [
+        ("Some_tablE", "SE"),
+        ("SomeTab_le", "ST"),
+    ],
+)
+def test_generate_alias_prefers_upper_case_name_over_underscore_name(table_name, alias):
+    assert pgcompleter.generate_alias(table_name) == alias