Merge pull request #624 from dbcli/koljonen/search_by_alias

Search suggestions by initialism

2 files changed, 151 insertions, 35 deletions

diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py
index f13e3f33..f7056655 100644
--- a/pgcli/pgcompleter.py
+++ b/pgcli/pgcompleter.py
@@ -32,17 +32,16 @@ NamedQueries.instance = NamedQueries.from_config(
 Match = namedtuple('Match', ['completion', 'priority'])
 
 _Candidate = namedtuple('Candidate', ['completion', 'priority', 'meta', 'synonyms'])
-def Candidate(completion, priority, meta, synonyms = None):
+def Candidate(completion, priority=None, meta=None, synonyms=None):
     return _Candidate(completion, priority, meta, synonyms or [completion])
 
 normalize_ref = lambda ref: ref if ref[0] == '"' else '"' + ref.lower() +  '"'
 
-def generate_alias(tbl, tbs):
+def generate_alias(tbl):
     """ 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
-    param tbls - set TableReference objects for tables already in query
     """
     return ''.join([l for l in tbl if l.isupper()] or
         [l for l, prev in zip(tbl,  '_' + tbl) if prev == '_' and l != '_'])
@@ -294,6 +293,12 @@ class PGCompleter(Completer):
             pat = re.compile('(%s)' % regex)
 
             def _match(item):
+                if item.lower()[:len(text) + 1] in (text, text + ' '):
+                    # Exact match of first word in suggestion
+                    # This is to get exact alias matches to the top
+                    # E.g. for input `e`, 'Entries E' should be on top
+                    # (before e.g. `EndUsers EU`)
+                    return float('Infinity'), -1
                 r = pat.search(self.unescape_name(item.lower()))
                 if r:
                     return -len(r.group()), -r.start()
@@ -311,7 +316,12 @@ class PGCompleter(Completer):
         for cand in collection:
             if isinstance(cand, _Candidate):
                 item, prio, display_meta, synonyms = cand
-                sort_key = max(_match(x) for x in synonyms)
+                if display_meta is None:
+                    display_meta = meta
+                syn_matches = (_match(x) for x in synonyms)
+                # Nones need to be removed to avoid max() crashing in Python 3
+                syn_matches = [m for m in syn_matches if m]
+                sort_key = max(syn_matches) if syn_matches else None
             else:
                 item, display_meta, prio = cand, meta, 0
                 sort_key = _match(cand)
@@ -390,7 +400,8 @@ class PGCompleter(Completer):
 
         colit = scoped_cols.items
         def make_cand(name, ref):
-            return Candidate(qualify(name, ref), 0, 'column', [name])
+            synonyms = (name, generate_alias(self.case(name)))
+            return Candidate(qualify(name, ref), 0, 'column', synonyms)
         flat_cols = []
         for t, cols in colit():
             for c in cols:
@@ -432,7 +443,7 @@ class PGCompleter(Completer):
         tbl = self.case(tbl)
         tbls = set(normalize_ref(t.ref) for t in tbls)
         if self.generate_aliases:
-            tbl = generate_alias(self.unescape_name(tbl), tbls)
+            tbl = generate_alias(self.unescape_name(tbl))
         if normalize_ref(tbl) not in tbls:
             return tbl
         elif tbl[0] == '"':
@@ -470,6 +481,9 @@ class PGCompleter(Completer):
             else:
                 join = '{0} ON {0}.{1} = {2}.{3}'.format(
                     c(left.tbl), c(left.col), rtbl.ref, c(right.col))
+            alias = generate_alias(self.case(left.tbl))
+            synonyms = [join, '{0} ON {0}.{1} = {2}.{3}'.format(
+                alias, c(left.col), rtbl.ref, c(right.col))]
             # Schema-qualify if (1) new table in same schema as old, and old
             # is schema-qualified, or (2) new in other schema, except public
             if not suggestion.schema and (qualified[normalize_ref(rtbl.ref)]
@@ -478,7 +492,7 @@ class PGCompleter(Completer):
                 join = left.schema + '.' + join
             prio = ref_prio[normalize_ref(rtbl.ref)] * 2 + (
                 0 if (left.schema, left.tbl) in other_tbls else 1)
-            joins.append(Candidate(join, prio, 'join'))
+            joins.append(Candidate(join, prio, 'join', synonyms=synonyms))
 
         return self.find_matches(word_before_cursor, joins, meta='join')
 
@@ -536,18 +550,16 @@ class PGCompleter(Completer):
         return self.find_matches(word_before_cursor, conds, meta='join')
 
     def get_function_matches(self, suggestion, word_before_cursor, alias=False):
+        def _cand(func_name, alias):
+            return self._make_cand(func_name, alias, suggestion, function=True)
         if suggestion.filter == 'for_from_clause':
             # Only suggest functions allowed in FROM clause
             filt = lambda f: not f.is_aggregate and not f.is_window
-            funcs = self.populate_functions(suggestion.schema, filt)
-            if alias:
-                funcs = [self.case(f) + '() ' + self.alias(f,
-                    suggestion.table_refs) for f in funcs]
-            else:
-                funcs = [self.case(f) + '()' for f in funcs]
+            funcs = [_cand(f, alias)
+                     for f in self.populate_functions(suggestion.schema, filt)]
         else:
-            funcs = [f + '()' for f in self.populate_schema_objects(
-                suggestion.schema, 'functions')]
+            fs = self.populate_schema_objects(suggestion.schema, 'functions')
+            funcs = [_cand(f, alias=False) for f in fs]
 
         # Function overloading means we way have multiple functions of the same
         # name at this point, so keep unique names only
@@ -585,6 +597,15 @@ class PGCompleter(Completer):
             + self.get_view_matches(v_sug, word_before_cursor, alias)
             + self.get_function_matches(f_sug, word_before_cursor, alias))
 
+    def _make_cand(self, tbl, do_alias, suggestion, function=False):
+        cased_tbl = self.case(tbl)
+        alias = self.alias(cased_tbl, suggestion.table_refs)
+        synonyms = (cased_tbl, generate_alias(cased_tbl))
+        maybe_parens = '()' if function else ''
+        maybe_alias = (' ' + alias) if do_alias else ''
+        item = cased_tbl + maybe_parens + maybe_alias
+        return Candidate(item, synonyms=synonyms)
+
     def get_table_matches(self, suggestion, word_before_cursor, alias=False):
         tables = self.populate_schema_objects(suggestion.schema, 'tables')
         tables.extend(tbl.name for tbl in suggestion.local_tables)
@@ -594,11 +615,8 @@ class PGCompleter(Completer):
         if not suggestion.schema and (
                 not word_before_cursor.startswith('pg_')):
             tables = [t for t in tables if not t.startswith('pg_')]
-        if alias:
-            tables = [self.case(t) + ' ' + self.alias(t, suggestion.table_refs)
-                for t in tables]
-        return self.find_matches(word_before_cursor, tables,
-            meta='table')
+        tables = [self._make_cand(t, alias, suggestion) for t in tables]
+        return self.find_matches(word_before_cursor, tables, meta='table')
 
 
     def get_view_matches(self, suggestion, word_before_cursor, alias=False):
@@ -607,9 +625,7 @@ class PGCompleter(Completer):
         if not suggestion.schema and (
                 not word_before_cursor.startswith('pg_')):
             views = [v for v in views if not v.startswith('pg_')]
-        if alias:
-            views = [self.case(v) + ' ' + self.alias(v, suggestion.table_refs)
-                for v in views]
+        views = [self._make_cand(v, alias, suggestion) for v in views]
         return self.find_matches(word_before_cursor, views, meta='view')
 
     def get_alias_matches(self, suggestion, word_before_cursor):
diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
index 608f3b84..2d2d599a 100644
--- a/tests/test_smart_completion_multiple_schemata.py
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -20,7 +20,14 @@ metadata = {
         },
         'Custom': {
             'projects': ['projectid', 'name']
-        }},
+        },
+        'blog': {
+            'entries': ['entryid', 'entrytitle', 'entrytext'],
+            'tags': ['tagid', 'name'],
+            'entrytags': ['entryid', 'tagid'],
+            'entacclog': ['entryid', 'username', 'datestamp'],
+        }
+    },
     'functions': {
         'public': [
             ['func1', [], [], [], '', False, False, False],
@@ -30,7 +37,13 @@ metadata = {
             ['set_returning_func', ['x'], ['integer'], ['o'],
                 'integer', False, False, True]],
         'Custom': [
-            ['func4', [], [], [], '', False, False, False]]
+            ['func4', [], [], [], '', False, False, False]],
+        'blog': [
+            ['extract_entry_symbols', ['_entryid', 'symbol'],
+                ['integer', 'text'], ['i', 'o'], '', False, False, True],
+            ['enter_entry', ['_title', '_text', 'entryid'],
+                ['text', 'text', 'integer'], ['i', 'i', 'o'],
+                '', False, False, False]],
      },
     'datatypes': {
         'public': ['typ1', 'typ2'],
@@ -39,16 +52,25 @@ metadata = {
     'foreignkeys': {
         'custom': [
             ('public', 'users', 'id', 'custom', 'shipments', 'user_id')
-    ]},
+        ],
+        'blog': [
+            ('blog', 'entries', 'entryid', 'blog', 'entacclog', 'entryid'),
+            ('blog', 'entries', 'entryid', 'blog', 'entrytags', 'entryid'),
+            ('blog', 'tags', 'tagid', 'blog', 'entrytags', 'tagid'),
+        ],
+    },
 }
 
 testdata = MetaData(metadata)
+cased_schemas = [schema(x) for x in ('public', 'blog', 'CUSTOM', '"Custom"')]
 
 @pytest.fixture
 def completer():
     return testdata.completer
 
-casing = ('SELECT', 'Orders', 'User_Emails', 'CUSTOM', 'Func1')
+casing = ('SELECT', 'Orders', 'User_Emails', 'CUSTOM', 'Func1', 'Entries',
+          'Tags', 'EntryTags', 'EntAccLog',
+          'EntryID', 'EntryTitle', 'EntryText')
 
 @pytest.fixture
 def completer_with_casing():
@@ -475,10 +497,7 @@ def test_table_aliases(completer_with_aliases, complete_event, text):
 def test_aliases_with_casing(completer_aliases_casing, complete_event, text):
     result = completer_aliases_casing.get_completions(
         Document(text=text), complete_event)
-    assert set(result) == set([
-        schema('public'),
-        schema('CUSTOM'),
-        schema('"Custom"'),
+    assert set(result) == set(cased_schemas + [
         table('users u'),
         table('Orders O' if text == 'SELECT * FROM ' else 'Orders O2'),
         table('"select" s'),
@@ -489,12 +508,93 @@ def test_aliases_with_casing(completer_aliases_casing, complete_event, text):
 def test_table_casing(completer_with_casing, complete_event, text):
     result = completer_with_casing.get_completions(
         Document(text=text), complete_event)
-    assert set(result) == set([
-        schema('public'),
-        schema('CUSTOM'),
-        schema('"Custom"'),
+    assert set(result) == set(cased_schemas + [
         table('users'),
         table('Orders'),
         table('"select"'),
         function('Func1()'),
         function('func2()')])
+
+def test_alias_search_without_aliases2(completer_with_casing, complete_event):
+    text = 'SELECT * FROM blog.et'
+    result = completer_with_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[0] == table('EntryTags', -2)
+
+def test_alias_search_without_aliases1(completer_with_casing, complete_event):
+    text = 'SELECT * FROM blog.e'
+    result = completer_with_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[0] == table('Entries', -1)
+
+def test_alias_search_with_aliases2(completer_aliases_casing, complete_event):
+    text = 'SELECT * FROM blog.et'
+    result = completer_aliases_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[0] == table('EntryTags ET', -2)
+
+def test_alias_search_with_aliases1(completer_aliases_casing, complete_event):
+    text = 'SELECT * FROM blog.e'
+    result = completer_aliases_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[0] == table('Entries E', -1)
+
+def test_join_alias_search_with_aliases1(completer_aliases_casing,
+                                         complete_event):
+    text = 'SELECT * FROM blog.Entries E JOIN blog.e'
+    result = completer_aliases_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[:2] == [table('Entries E2', -1), join(
+        'EntAccLog EAL ON EAL.EntryID = E.EntryID', -1)]
+
+def test_join_alias_search_without_aliases1(completer_with_casing,
+                                            complete_event):
+    text = 'SELECT * FROM blog.Entries JOIN blog.e'
+    result = completer_with_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[:2] == [table('Entries', -1), join(
+        'EntAccLog ON EntAccLog.EntryID = Entries.EntryID', -1)]
+
+def test_join_alias_search_with_aliases2(completer_aliases_casing,
+                                         complete_event):
+    text = 'SELECT * FROM blog.Entries E JOIN blog.et'
+    result = completer_aliases_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[0] == join('EntryTags ET ON ET.EntryID = E.EntryID', -2)
+
+def test_join_alias_search_without_aliases2(completer_with_casing,
+                                            complete_event):
+    text = 'SELECT * FROM blog.Entries JOIN blog.et'
+    result = completer_with_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[0] == join(
+        'EntryTags ON EntryTags.EntryID = Entries.EntryID', -2)
+
+def test_function_alias_search_without_aliases(completer_with_casing,
+                                               complete_event):
+    text = 'SELECT blog.ees'
+    result = completer_with_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[0] == function('extract_entry_symbols()', -3)
+
+def test_function_alias_search_with_aliases(completer_aliases_casing,
+                                            complete_event):
+    text = 'SELECT blog.ee'
+    result = completer_aliases_casing.get_completions(
+        Document(text=text), complete_event)
+    assert result[0] == function('enter_entry()', -2)
+
+def test_column_alias_search(completer_aliases_casing, complete_event):
+    text = 'SELECT et FROM blog.Entries E'
+    result = completer_aliases_casing.get_completions(
+        Document(text, cursor_position=len('SELECT et')), complete_event)
+    cols = ('EntryText', 'EntryTitle', 'EntryID')
+    assert result[:3] == [column(c, -2) for c in cols]
+
+def test_column_alias_search_qualified(completer_aliases_casing,
+                                       complete_event):
+    text = 'SELECT E.ei FROM blog.Entries E'
+    result = completer_aliases_casing.get_completions(
+        Document(text, cursor_position=len('SELECT E.ei')), complete_event)
+    cols = ('EntryID', 'EntryTitle')
+    assert result[:3] == [column(c, -2) for c in cols]