Take account of table aliases when completing function args (#1048)

3 files changed, 50 insertions, 28 deletions

diff --git a/changelog.rst b/changelog.rst
index 1ec7b729..60404842 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -6,6 +6,7 @@ Features:
 
 * Add `\\G` as a terminator to sql statements that will show the results in expanded mode. This feature is copied from mycli. (Thanks: `Amjith Ramanujam`_)
 * Removed limit prompt and added automatic row limit on queries with no LIMIT clause (#1079) (Thanks: `Sebastian Janko`_)
+* Function argument completions now take account of table aliases (#1048). (Thanks: `Owen Stephens`_)
 
 Bug fixes:
 ----------
diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index 85b3bb10..781a24ab 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -398,35 +398,11 @@ def suggest_based_on_last_token(token, stmt):
         if prev_prev_tok and prev_prev_tok.normalized == "INTO":
             return (Column(table_refs=stmt.get_tables("insert"), context="insert"),)
         # We're probably in a function argument list
-        return (
-            Column(
-                table_refs=extract_tables(stmt.full_text),
-                local_tables=stmt.local_tables,
-                qualifiable=True,
-            ),
-        )
+        return _suggest_expression(token_v, stmt)
     elif token_v == "set":
         return (Column(table_refs=stmt.get_tables(), local_tables=stmt.local_tables),)
     elif token_v in ("select", "where", "having", "order by", "distinct"):
-        # Check for a table alias or schema qualification
-        parent = (stmt.identifier and stmt.identifier.get_parent_name()) or []
-        tables = stmt.get_tables()
-        if parent:
-            tables = tuple(t for t in tables if identifies(parent, t))
-            return (
-                Column(table_refs=tables, local_tables=stmt.local_tables),
-                Table(schema=parent),
-                View(schema=parent),
-                Function(schema=parent),
-            )
-        else:
-            return (
-                Column(
-                    table_refs=tables, local_tables=stmt.local_tables, qualifiable=True
-                ),
-                Function(schema=None),
-                Keyword(token_v.upper()),
-            )
+        return _suggest_expression(token_v, stmt)
     elif token_v == "as":
         # Don't suggest anything for aliases
         return ()
@@ -554,6 +530,30 @@ def suggest_based_on_last_token(token, stmt):
         return (Keyword(),)
 
 
+def _suggest_expression(token_v, stmt):
+    """
+    Return suggestions for an expression, taking account of the partially-typed
+    identifiers parent, which may be a table alias or schema name.
+    """
+    parent = stmt.identifier.get_parent_name() if stmt.identifier else []
+    tables = stmt.get_tables()
+
+    if parent:
+        tables = tuple(t for t in tables if identifies(parent, t))
+        return (
+            Column(table_refs=tables, local_tables=stmt.local_tables),
+            Table(schema=parent),
+            View(schema=parent),
+            Function(schema=parent),
+        )
+
+    return (
+        Column(table_refs=tables, local_tables=stmt.local_tables, qualifiable=True),
+        Function(schema=None),
+        Keyword(token_v.upper()),
+    )
+
+
 def identifies(id, ref):
     """Returns true if string `id` matches TableReference `ref`"""
 
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index ed647fc6..c5ef7b7d 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -101,10 +101,14 @@ def test_where_equals_any_suggests_columns_or_keywords():
     assert set(suggestions) == cols_etc("tabl", last_keyword="WHERE")
 
 
-def test_lparen_suggests_cols():
+def test_lparen_suggests_cols_and_funcs():
     suggestion = suggest_type("SELECT MAX( FROM tbl", "SELECT MAX(")
     assert set(suggestion) == set(
-        [Column(table_refs=((None, "tbl", None, False),), qualifiable=True)]
+        [
+            Column(table_refs=((None, "tbl", None, False),), qualifiable=True),
+            Function(schema=None),
+            Keyword("("),
+        ]
     )
 
 
@@ -282,6 +286,23 @@ def test_distinct_and_order_by_suggestions_with_alias_given(text, text_before):
     )
 
 
+def test_function_arguments_with_alias_given():
+    suggestions = suggest_type("SELECT avg(x. FROM tbl x, tbl2 y", "SELECT avg(x.")
+
+    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(