Find statements inside function body

Consider this script
```
CREATE FUNCTION foo() returns text LANGUAGE SQL AS $func$
SELECT 1 FROM Bar;
SELECT <cursor> FROM Baz;
$func$;
```
The change here is that `SELECT <cursor> FROM Baz;` will be seen as the
current statement, instead of the whole function definition.
This means we'll no longer get column suggestions from `Bar`.

changelog.rst

@@ -1,3 +1,10 @@
+Upcoming
+=====
+
+Features:
+---------
+* Better suggestions when editing functions (Thanks: `Joakim Koljonen`_)
+
 1.5.0
 =====
 

pgcli/packages/sqlcompletion.py

@@ -168,6 +168,26 @@ def _strip_named_query(txt):
         txt = named_query_regex.sub('', txt)
     return txt
 
+function_body_pattern = re.compile('(\\$.*?\\$)([\s\S]*?)\\1', re.M)
+
+
+def _find_function_body(text):
+    split = function_body_pattern.search(text)
+    return (split.start(2), split.end(2)) if split else (None, None)
+
+
+def _statement_from_function(full_text, text_before_cursor, statement):
+    current_pos = len(text_before_cursor)
+    body_start, body_end = _find_function_body(full_text)
+    if body_start is None:
+        return full_text, text_before_cursor, statement
+    if not body_start <= current_pos < body_end:
+        return full_text, text_before_cursor, statement
+    full_text = full_text[body_start:body_end]
+    text_before_cursor = text_before_cursor[body_start:]
+    parsed = sqlparse.parse(text_before_cursor)
+    return _split_multiple_statements(full_text, text_before_cursor, parsed)
+
 
 def _split_multiple_statements(full_text, text_before_cursor, parsed):
     if len(parsed) > 1:
@@ -191,8 +211,18 @@ def _split_multiple_statements(full_text, text_before_cursor, parsed):
         statement = parsed[0]
     else:
         # The empty string
-        statement = None
+        return full_text, text_before_cursor, None
 
+    token2 = None
+    if statement.get_type() in ('CREATE', 'CREATE OR REPLACE'):
+        token1 = statement.token_first()
+        if token1:
+            token1_idx = statement.token_index(token1)
+            token2 = statement.token_next(token1_idx)[1]
+    if token2 and token2.value.upper() == 'FUNCTION':
+        full_text, text_before_cursor, statement = _statement_from_function(
+            full_text, text_before_cursor, statement
+        )
     return full_text, text_before_cursor, statement
 
 

tests/test_smart_completion_multiple_schemata.py

@@ -387,6 +387,17 @@ def test_wildcard_column_expansion_with_alias_qualifier(completer, complete_even
     assert expected == completions
 
 @pytest.mark.parametrize('text', [
+    '''
+    SELECT count(1) FROM users;
+    CREATE FUNCTION foo(custom.products _products) returns custom.shipments
+    LANGUAGE SQL
+    AS $foo$
+    SELECT 1 FROM custom.shipments;
+    INSERT INTO public.orders(*) values(-1, now(), 'preliminary');
+    SELECT 2 FROM custom.users;
+    $foo$;
+    SELECT count(1) FROM custom.shipments;
+    ''',
     'INSERT INTO public.orders(*',
     'INSERT INTO public.Orders(*',
     'INSERT INTO public.orders (*',

tests/test_sqlcompletion.py

@@ -582,6 +582,85 @@ def test_3_statements_2nd_current():
                                'select * from a; select ')
     assert set(suggestions) == cols_etc('b')
 
+@pytest.mark.parametrize('text', [
+'''
+CREATE OR REPLACE FUNCTION func() RETURNS setof int AS $$
+SELECT  FROM foo;
+SELECT 2 FROM bar;
+$$ language sql;
+    ''',
+    '''create function func2(int, varchar)
+RETURNS text
+language sql AS
+$func$
+SELECT 2 FROM bar;
+SELECT  FROM foo;
+$func$
+    ''',
+'''
+CREATE OR REPLACE FUNCTION func() RETURNS setof int AS $func$
+SELECT 3 FROM foo;
+SELECT 2 FROM bar;
+$$ language sql;
+create function func2(int, varchar)
+RETURNS text
+language sql AS
+$func$
+SELECT 2 FROM bar;
+SELECT  FROM foo;
+$func$
+    ''',
+'''
+SELECT * FROM baz;
+CREATE OR REPLACE FUNCTION func() RETURNS setof int AS $func$
+SELECT  FROM foo;
+SELECT 2 FROM bar;
+$$ language sql;
+create function func2(int, varchar)
+RETURNS text
+language sql AS
+$func$
+SELECT 3 FROM bar;
+SELECT  FROM foo;
+$func$
+SELECT * FROM qux;
+    '''
+])
+def test_statements_in_function_body(text):
+    suggestions = suggest_type(text, text[:text.find('  ') + 1])
+    assert set(suggestions) == set([
+        Column(table_refs=((None, 'foo', None, False),), qualifiable=True),
+        Function(schema=None),
+        Keyword()
+    ])
+
+functions = [
+'''
+CREATE OR REPLACE FUNCTION func() RETURNS setof int AS $$
+SELECT 1 FROM foo;
+SELECT 2 FROM bar;
+$$ language sql;
+    ''',
+    '''
+create function func2(int, varchar)
+RETURNS text
+language sql AS
+'
+SELECT 2 FROM bar;
+SELECT 1 FROM foo;
+';
+    '''
+]
+
+@pytest.mark.parametrize('text', functions)
+def test_statements_with_cursor_after_function_body(text):
+    suggestions = suggest_type(text, text[:text.find('; ') + 1])
+    assert set(suggestions) == set([Keyword()])
+
+@pytest.mark.parametrize('text', functions)
+def test_statements_with_cursor_before_function_body(text):
+    suggestions = suggest_type(text, '')
+    assert set(suggestions) == set([Keyword()])
 
 def test_create_db_with_template():
     suggestions = suggest_type('create database foo with template ',