Add initial completion support for JOIN statements.

3 files changed, 34 insertions, 4 deletions

diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index b8c94f88..d447fefc 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -1,5 +1,6 @@
 from __future__ import print_function
 import sqlparse
+from sqlparse.sql import Comparison
 from parseutils import last_word, extract_tables, find_prev_keyword
 
 
@@ -39,7 +40,16 @@ def suggest_based_on_last_token(token, text_before_cursor, full_text):
     if isinstance(token, basestring):
         token_v = token
     else:
-        token_v = token.value
+        # If 'token' is a Comparison type such as
+        # 'select * FROM abc a JOIN def d ON a.id = d.'. Then calling
+        # token.value on the comparison type will only return the lhs of the
+        # comparison. In this case a.id. So we need to do token.tokens to get
+        # both sides of the comparison and pick the last token out of that
+        # list.
+        if isinstance(token, Comparison):
+            token_v = token.tokens[-1].value
+        else:
+            token_v = token.value
 
     if token_v.lower().endswith('('):
         p = sqlparse.parse(text_before_cursor)[0]
@@ -53,7 +63,7 @@ def suggest_based_on_last_token(token, text_before_cursor, full_text):
         return 'columns', extract_tables(full_text)
     elif token_v.lower() in ('select', 'where', 'having'):
         return 'columns-and-functions', extract_tables(full_text)
-    elif token_v.lower() in ('from', 'update', 'into', 'describe'):
+    elif token_v.lower() in ('from', 'update', 'into', 'describe', 'join'):
         return 'tables', []
     elif token_v in ('d',):  # \d
         return 'tables', []
diff --git a/tests/test_parseutils.py b/tests/test_parseutils.py
index ae6e8f5a..bc85b58c 100644
--- a/tests/test_parseutils.py
+++ b/tests/test_parseutils.py
@@ -38,8 +38,12 @@ def test_simple_update_table():
     assert tables == ['abc']
 
 def test_join_table():
-    tables = extract_tables('SELECT * FROM abc a JOIN def d ON s.id = a.num')
-    assert tables == ['abc', 'def']
+    expected = {'a': 'abc', 'd': 'def'}
+    tables = extract_tables('SELECT * FROM abc a JOIN def d ON a.id = d.num')
+    tables_aliases = extract_tables(
+            'SELECT * FROM abc a JOIN def d ON a.id = d.num', True)
+    assert tables == expected.values()
+    assert tables_aliases == expected
 
 def test_join_as_table():
     expected = {'m': 'my_table'}
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index 86ae896c..3b76653b 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -66,6 +66,7 @@ def test_dot_suggests_cols_of_an_alias():
     assert suggestion == ('columns', ['tabl1'])
 
 def test_dot_col_comma_suggests_cols():
+    import pdb; pdb.set_trace()
     suggestion = suggest_type('SELECT t1.a, t2. FROM tabl1 t1, tabl2 t2',
             'SELECT t1.a, t2.')
     assert suggestion == ('columns', ['tabl2'])
@@ -97,3 +98,18 @@ def test_sub_select_dot_col_name_completion():
     suggestion = suggest_type('SELECT * FROM (SELECT t. FROM tabl t',
             'SELECT * FROM (SELECT t.')
     assert suggestion == ('columns', ['tabl'])
+
+def test_join_suggests_tables():
+    suggestion = suggest_type('SELECT * FROM abc a JOIN ',
+            'SELECT * FROM abc a JOIN ')
+    assert suggestion == ('tables', [])
+
+def test_join_alias_dot_suggests_cols1():
+    suggestion = suggest_type('SELECT * FROM abc a JOIN def d ON a.',
+            'SELECT * FROM abc a JOIN def d ON a.')
+    assert suggestion == ('columns', ['abc'])
+
+def test_join_alias_dot_suggests_cols2():
+    suggestion = suggest_type('SELECT * FROM abc a JOIN def d ON a.',
+            'SELECT * FROM abc a JOIN def d ON a.id = d.')
+    assert suggestion == ('columns', ['def'])