Suggest columns for `ORDER BY` and `DISTINCT` (fixes #685)

Having typed an alias name in an `ORDER BY` or (`SELECT`) `DISTINCT`
clause, the alias was not taken account of, and the completion simply
listed all columns. This change fixes that, and makes the autocompletion
behave the same as in `SELECT` and `WHERE` clauses.

3 files changed, 58 insertions, 6 deletions

diff --git a/changelog.rst b/changelog.rst
index 96d4893f..45856a25 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -5,6 +5,7 @@ Features:
 ---------
 * Suggest objects from all schemas (not just those in search_path) (Thanks: `Joakim Koljonen`_)
 * Allow configurable character to be used for multi-line query continuations. (Thanks: `Owen Stephens`_)
+* Completions after ORDER BY and DISTINCT now take account of table aliases. (Thanks: `Owen Stephens`_)
 
 Bug fixes:
 ----------
diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index a870d1e1..cd5e151a 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -376,10 +376,7 @@ def suggest_based_on_last_token(token, stmt):
     elif token_v == 'set':
         return (Column(table_refs=stmt.get_tables(),
                        local_tables=stmt.local_tables),)
-    elif token_v in ('by', 'distinct'):
-        return (Column(table_refs=stmt.get_tables(),
-                       local_tables=stmt.local_tables, qualifiable=True),)
-    elif token_v in ('select', 'where', 'having'):
+    elif token_v in ('select', 'where', 'having', 'by', 'distinct'):
         # Check for a table alias or schema qualification
         parent = (stmt.identifier and stmt.identifier.get_parent_name()) or []
         tables = stmt.get_tables()
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index a2928df4..c4a7da11 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -4,8 +4,9 @@ from pgcli.packages.sqlcompletion import (
 from pgcli.packages.parseutils.tables import TableReference
 import pytest
 
-# Returns the expected select-clause suggestions for a single-table select
 def cols_etc(table, schema=None, alias=None, is_function=False, parent=None):
+    """Returns the expected select-clause suggestions for a single-table
+    select."""
     return set([
         Column(table_refs=(TableReference(schema, table, alias, is_function),),
                qualifiable=True),
@@ -213,9 +214,62 @@ def test_truncate_suggests_qualified_tables():
 ])
 def test_distinct_suggests_cols(text):
     suggestions = suggest_type(text, text)
-    assert suggestions ==(Column(table_refs=(), qualifiable=True),)
+    assert set(suggestions) == set([
+        Column(table_refs=(), local_tables=(), qualifiable=True),
+        Function(schema=None),
+        Keyword()
+    ])
 
 
+@pytest.mark.parametrize('text, text_before', [
+    (
+        'SELECT DISTINCT FROM tbl x JOIN tbl1 y',
+        'SELECT DISTINCT'
+    ),
+    (
+        'SELECT * FROM tbl x JOIN tbl1 y ORDER BY ',
+        'SELECT * FROM tbl x JOIN tbl1 y ORDER BY '
+    )
+])
+def test_distinct_and_order_by_suggestions_with_aliases(text, text_before):
+    suggestions = suggest_type(text, text_before)
+    assert set(suggestions) == set([
+        Column(
+            table_refs=(
+                TableReference(None, 'tbl', 'x', False),
+                TableReference(None, 'tbl1', 'y', False),
+            ),
+            local_tables=(),
+            qualifiable=True
+        ),
+        Function(schema=None),
+        Keyword()
+    ])
+
+
+@pytest.mark.parametrize('text, text_before', [
+    (
+        'SELECT DISTINCT x. FROM tbl x JOIN tbl1 y',
+        'SELECT DISTINCT x.'
+    ),
+    (
+        'SELECT * FROM tbl x JOIN tbl1 y ORDER BY x.',
+        'SELECT * FROM tbl x JOIN tbl1 y ORDER BY x.'
+    )
+])
+def test_distinct_and_order_by_suggestions_with_alias_given(text, text_before):
+    suggestions = suggest_type(text, text_before)
+    assert set(suggestions) == set([
+        Column(
+            table_refs=(TableReference(None, 'tbl', 'x', False),),
+            local_tables=(),
+            qualifiable=False
+        ),
+        Table(schema='x'),
+        View(schema='x'),
+        Function(schema='x'),
+    ])
+
 def test_col_comma_suggests_cols():
     suggestions = suggest_type('SELECT a, b, FROM tbl', 'SELECT a, b,')
     assert set(suggestions) == set([