Merge pull request #142 from darikg/master

Fix #106 - autocompletion in multiple statements

2 files changed, 67 insertions, 4 deletions

diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index 45e38fed..75d9f98d 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -38,10 +38,30 @@ def suggest_type(full_text, text_before_cursor):
     else:
         parsed = sqlparse.parse(text_before_cursor)
 
-    # Need to check if `p` is not empty, since an empty string will result in
-    # an empty tuple.
-    p = parsed[0] if parsed else None
-    last_token = p and p.token_prev(len(p.tokens)) or ''
+    if len(parsed) > 1:
+        # Multiple statements being edited -- isolate the current one by
+        # cumulatively summing statement lengths to find the one that bounds the
+        # current position
+        current_pos = len(text_before_cursor)
+        stmt_start, stmt_end = 0, 0
+
+        for statement in parsed:
+            stmt_len = len(statement.to_unicode())
+            stmt_start, stmt_end = stmt_end, stmt_end + stmt_len
+
+            if stmt_end >= current_pos:
+                text_before_cursor = full_text[stmt_start:current_pos]
+                full_text = full_text[stmt_start:]
+                break
+
+    elif parsed:
+        # A single statement
+        statement = parsed[0]
+    else:
+        # The empty string
+        statement = None
+
+    last_token = statement and statement.token_prev(len(statement.tokens)) or ''
 
     return suggest_based_on_last_token(last_token, text_before_cursor, full_text)
 
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index 22e8c11b..760e4dba 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -183,3 +183,46 @@ def test_on_suggests_tables_right_side():
         'select abc.x, bcd.y from abc join bcd on ',
         'select abc.x, bcd.y from abc join bcd on ')
     assert suggestions == [{'type': 'alias', 'aliases': ['abc', 'bcd']}]
+
+def test_2_statements_2nd_current():
+    suggestions = suggest_type('select * from a; select * from ',
+                               'select * from a; select * from ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+         {'type': 'table', 'schema': []}, {'type': 'schema'}])
+
+    suggestions = suggest_type('select * from a; select  from b',
+                               'select * from a; select ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+            {'type': 'column', 'tables': [(None, 'b', None)]},
+            {'type': 'function'}])
+
+    # Should work even if first statement is invalid
+    suggestions = suggest_type('select * from; select * from ',
+                               'select * from; select * from ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+         {'type': 'table', 'schema': []}, {'type': 'schema'}])
+
+def test_2_statements_1st_current():
+    suggestions = suggest_type('select * from ; select * from b',
+                               'select * from ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+         {'type': 'table', 'schema': []}, {'type': 'schema'}])
+
+    suggestions = suggest_type('select  from a; select * from b',
+                               'select ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+            {'type': 'column', 'tables': [(None, 'a', None)]},
+            {'type': 'function'}])
+
+def test_3_statements_2nd_current():
+    suggestions = suggest_type('select * from a; select * from ; select * from c',
+                               'select * from a; select * from ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+         {'type': 'table', 'schema': []}, {'type': 'schema'}])
+
+    suggestions = suggest_type('select * from a; select  from b; select * from c',
+                               'select * from a; select ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+            {'type': 'column', 'tables': [(None, 'b', None)]},
+            {'type': 'function'}])
+