Merge pull request #1107 from owst/function_arg_completions_with_alias

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

5 files changed, 54 insertions, 30 deletions

diff --git a/changelog.rst b/changelog.rst
index 225707c6..0d7a0ba0 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..2828b175 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 any partially-typed
+    identifier's 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_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
index 9125f124..9fff2a17 100644
--- a/tests/test_smart_completion_multiple_schemata.py
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -229,7 +229,7 @@ def test_suggested_column_names_in_function(completer):
         completer, "SELECT MAX( from custom.products", len("SELECT MAX(")
     )
     assert completions_to_set(result) == completions_to_set(
-        testdata.columns("products", "custom")
+        testdata.columns_functions_and_keywords("products", "custom")
     )
 
 
diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py
index 63f1e80e..bb371543 100644
--- a/tests/test_smart_completion_public_schema_only.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -287,7 +287,9 @@ def test_suggested_cased_always_qualified_column_names(completer):
 @parametrize("completer", completers(casing=False, qualify=no_qual))
 def test_suggested_column_names_in_function(completer):
     result = get_result(completer, "SELECT MAX( from users", len("SELECT MAX("))
-    assert completions_to_set(result) == completions_to_set(testdata.columns("users"))
+    assert completions_to_set(result) == completions_to_set(
+        (testdata.columns_functions_and_keywords("users"))
+    )
 
 
 @parametrize("completer", completers(casing=False))
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(