ON keyword now suggests tables and aliases on the right side of equal sign also.

3 files changed, 35 insertions, 1 deletions

diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index b506a623..f409dae0 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -81,7 +81,7 @@ def suggest_based_on_last_token(token, text_before_cursor, full_text):
         return 'tables', []
     elif token_v.lower() in ('c', 'use'):  # \c
         return 'databases', []
-    elif token_v.endswith(','):
+    elif token_v.endswith(',') or token_v == '=':
         prev_keyword = find_prev_keyword(text_before_cursor)
         if prev_keyword:
             return suggest_based_on_last_token(prev_keyword, text_before_cursor, full_text)
diff --git a/tests/test_smart_completion.py b/tests/test_smart_completion.py
index b306b36a..425344b9 100644
--- a/tests/test_smart_completion.py
+++ b/tests/test_smart_completion.py
@@ -197,6 +197,16 @@ def test_suggested_aliases_after_on(completer, complete_event):
         Completion(text='u', start_position=0),
         Completion(text='o', start_position=0)])
 
+def test_suggested_aliases_after_on_right_side(completer, complete_event):
+    text = 'SELECT u.name, o.id FROM users u JOIN orders o ON o.user_id = '
+    position = len('SELECT u.name, o.id FROM users u JOIN orders o ON o.user_id = ')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='u', start_position=0),
+        Completion(text='o', start_position=0)])
+
 def test_suggested_tables_after_on(completer, complete_event):
     text = 'SELECT users.name, orders.id FROM users JOIN orders ON '
     position = len('SELECT users.name, orders.id FROM users JOIN orders ON ')
@@ -207,6 +217,16 @@ def test_suggested_tables_after_on(completer, complete_event):
         Completion(text='users', start_position=0),
         Completion(text='orders', start_position=0)])
 
+def test_suggested_tables_after_on_right_side(completer, complete_event):
+    text = 'SELECT users.name, orders.id FROM users JOIN orders ON orders.user_id = '
+    position = len('SELECT users.name, orders.id FROM users JOIN orders ON orders.user_id = ')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='users', start_position=0),
+        Completion(text='orders', start_position=0)])
+
 def test_table_names_after_from(completer, complete_event):
     text = 'SELECT * FROM '
     position = len('SELECT * FROM ')
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index 949321e7..df0cdc7c 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -127,3 +127,17 @@ def test_on_suggests_tables():
         'select abc.x, bcd.y from abc join bcd on ')
     assert category == 'tables-or-aliases'
     assert set(scope) == set(['abc', 'bcd'])
+
+def test_on_suggests_aliases_right_side():
+    category, scope = suggest_type(
+        'select a.x, b.y from abc a join bcd b on a.id = ',
+        'select a.x, b.y from abc a join bcd b on a.id = ')
+    assert category == 'tables-or-aliases'
+    assert set(scope) == set(['a', 'b'])
+
+def test_on_suggests_tables_right_side():
+    category, scope = suggest_type(
+        'select abc.x, bcd.y from abc join bcd on ',
+        'select abc.x, bcd.y from abc join bcd on ')
+    assert category == 'tables-or-aliases'
+    assert set(scope) == set(['abc', 'bcd'])