Merge pull request #553 from dbcli/darikg/cte-suggestions

CTE-aware suggestions

7 files changed, 206 insertions, 38 deletions

diff --git a/tests/metadata.py b/tests/metadata.py
index f495c14c..6630bbbb 100644
--- a/tests/metadata.py
+++ b/tests/metadata.py
@@ -1,4 +1,4 @@
-from pgcli.packages.function_metadata import FunctionMetadata, ForeignKey
+from pgcli.packages.parseutils.meta import FunctionMetadata, ForeignKey
 from prompt_toolkit.completion import Completion
 from functools import partial
 
diff --git a/tests/parseutils/test_ctes.py b/tests/parseutils/test_ctes.py
new file mode 100644
index 00000000..9566cf65
--- /dev/null
+++ b/tests/parseutils/test_ctes.py
@@ -0,0 +1,121 @@
+import pytest
+from sqlparse import parse
+from pgcli.packages.parseutils.ctes import (
+    token_start_pos, extract_ctes,
+    extract_column_names as _extract_column_names)
+
+
+def extract_column_names(sql):
+    p = parse(sql)[0]
+    return _extract_column_names(p)
+
+
+def test_token_str_pos():
+    sql = 'SELECT * FROM xxx'
+    p = parse(sql)[0]
+    idx = p.token_index(p.tokens[-1])
+    assert token_start_pos(p.tokens, idx) == len('SELECT * FROM ')
+
+    sql = 'SELECT * FROM \nxxx'
+    p = parse(sql)[0]
+    idx = p.token_index(p.tokens[-1])
+    assert token_start_pos(p.tokens, idx) == len('SELECT * FROM \n')
+
+
+def test_single_column_name_extraction():
+    sql = 'SELECT abc FROM xxx'
+    assert extract_column_names(sql) == ('abc',)
+
+
+def test_aliased_single_column_name_extraction():
+    sql = 'SELECT abc def FROM xxx'
+    assert extract_column_names(sql) == ('def',)
+
+
+def test_aliased_expression_name_extraction():
+    sql = 'SELECT 99 abc FROM xxx'
+    assert extract_column_names(sql) == ('abc',)
+
+
+def test_multiple_column_name_extraction():
+    sql = 'SELECT abc, def FROM xxx'
+    assert extract_column_names(sql) == ('abc', 'def')
+
+
+def test_missing_column_name_handled_gracefully():
+    sql = 'SELECT abc, 99 FROM xxx'
+    assert extract_column_names(sql) == ('abc',)
+
+    sql = 'SELECT abc, 99, def FROM xxx'
+    assert extract_column_names(sql) == ('abc', 'def')
+
+
+def test_aliased_multiple_column_name_extraction():
+    sql = 'SELECT abc def, ghi jkl FROM xxx'
+    assert extract_column_names(sql) == ('def', 'jkl')
+
+
+def test_table_qualified_column_name_extraction():
+    sql = 'SELECT abc.def, ghi.jkl FROM xxx'
+    assert extract_column_names(sql) == ('def', 'jkl')
+
+
+@pytest.mark.parametrize('sql', [
+    'INSERT INTO foo (x, y, z) VALUES (5, 6, 7) RETURNING x, y',
+    'DELETE FROM foo WHERE x > y RETURNING x, y',
+    'UPDATE foo SET x = 9 RETURNING x, y',
+])
+def test_extract_column_names_from_returning_clause(sql):
+    assert extract_column_names(sql) == ('x', 'y')
+
+
+def test_simple_cte_extraction():
+    sql = 'WITH a AS (SELECT abc FROM xxx) SELECT * FROM a'
+    start_pos = len('WITH a AS ')
+    stop_pos = len('WITH a AS (SELECT abc FROM xxx)')
+    ctes, remainder = extract_ctes(sql)
+
+    assert tuple(ctes) == (('a', ('abc',), start_pos, stop_pos),)
+    assert remainder.strip() == 'SELECT * FROM a'
+
+
+def test_cte_extraction_around_comments():
+    sql = '''--blah blah blah
+            WITH a AS (SELECT abc def FROM x)
+            SELECT * FROM a'''
+    start_pos = len('''--blah blah blah
+            WITH a AS ''')
+    stop_pos = len('''--blah blah blah
+            WITH a AS (SELECT abc def FROM x)''')
+
+    ctes, remainder = extract_ctes(sql)
+    assert tuple(ctes) == (('a', ('def',), start_pos, stop_pos),)
+    assert remainder.strip() == 'SELECT * FROM a'
+
+
+def test_multiple_cte_extraction():
+    sql = '''WITH
+            x AS (SELECT abc, def FROM x),
+            y AS (SELECT ghi, jkl FROM y)
+            SELECT * FROM a, b'''
+
+    start1 = len('''WITH
+            x AS ''')
+
+    stop1 = len('''WITH
+            x AS (SELECT abc, def FROM x)''')
+
+    start2 = len('''WITH
+            x AS (SELECT abc, def FROM x),
+            y AS ''')
+
+    stop2 = len('''WITH
+            x AS (SELECT abc, def FROM x),
+            y AS (SELECT ghi, jkl FROM y)''')
+
+    ctes, remainder = extract_ctes(sql)
+    assert tuple(ctes) == (
+        ('x', ('abc', 'def'), start1, stop1),
+        ('y', ('ghi', 'jkl'), start2, stop2))
+
+
diff --git a/tests/test_function_metadata.py b/tests/parseutils/test_function_metadata.py
index f7669e6f..097ce62d 100644
--- a/tests/test_function_metadata.py
+++ b/tests/parseutils/test_function_metadata.py
@@ -1,4 +1,4 @@
-from pgcli.packages.function_metadata import FunctionMetadata
+from pgcli.packages.parseutils.meta import FunctionMetadata
 
 
 def test_function_metadata_eq():
diff --git a/tests/test_parseutils.py b/tests/parseutils/test_parseutils.py
index f82f271f..6c2994f8 100644
--- a/tests/test_parseutils.py
+++ b/tests/parseutils/test_parseutils.py
@@ -1,6 +1,6 @@
 import pytest
-from pgcli.packages.parseutils import extract_tables
-from pgcli.packages.parseutils import find_prev_keyword, is_open_quote
+from pgcli.packages.parseutils.tables import extract_tables
+from pgcli.packages.parseutils.utils import find_prev_keyword, is_open_quote
 
 def test_empty_string():
     tables = extract_tables('')
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index e85cf36f..a9ae204f 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -3,7 +3,7 @@
 import pytest
 import psycopg2
 from pgspecial.main import PGSpecial
-from pgcli.packages.function_metadata import FunctionMetadata
+from pgcli.packages.parseutils.meta import FunctionMetadata
 from textwrap import dedent
 from utils import run, dbtest, requires_json, requires_jsonb
 
diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py
index 4802db45..c1d65ab5 100644
--- a/tests/test_smart_completion_public_schema_only.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -3,6 +3,7 @@ import pytest
 from metadata import (MetaData, alias, name_join, fk_join, join, keyword,
     schema, table, view, function, column, wildcard_expansion)
 from prompt_toolkit.document import Document
+from prompt_toolkit.completion import Completion
 
 metadata = {
     'tables': {
@@ -867,3 +868,49 @@ def test_insert(completer, complete_event, text):
     result = completer.get_completions(Document(text=text, cursor_position=pos),
                                        complete_event)
     assert set(result) == set(testdata.columns('users'))
+
+
+def test_suggest_cte_names(completer, complete_event):
+    text = '''
+        WITH cte1 AS (SELECT a, b, c FROM foo),
+             cte2 AS (SELECT d, e, f FROM bar)
+        SELECT * FROM
+    '''
+    pos = len(text)
+    result = completer.get_completions(
+        Document(text=text, cursor_position=pos),
+        complete_event)
+    expected = set([
+        Completion('cte1', 0, display_meta='table'),
+        Completion('cte2', 0, display_meta='table'),
+    ])
+    assert expected <= set(result)
+
+
+def test_suggest_columns_from_cte(completer, complete_event):
+    text = 'WITH cte AS (SELECT foo, bar FROM baz) SELECT  FROM cte'
+    pos = len('WITH cte AS (SELECT foo, bar FROM baz) SELECT ')
+    result = completer.get_completions(Document(text=text, cursor_position=pos),
+                                       complete_event)
+    expected = ([Completion('foo', 0, display_meta='column'),
+                 Completion('bar', 0, display_meta='column'),
+                 ] +
+                testdata.functions() +
+                testdata.builtin_functions() +
+                testdata.keywords()
+                )
+
+    assert set(expected) == set(result)
+
+
+@pytest.mark.parametrize('text', [
+    'WITH cte AS (SELECT foo FROM bar) SELECT * FROM cte WHERE cte.',
+    'WITH cte AS (SELECT foo FROM bar) SELECT * FROM cte c WHERE c.',
+])
+def test_cte_qualified_columns(completer, complete_event, text):
+    pos = len(text)
+    result = completer.get_completions(
+        Document(text=text, cursor_position=pos),
+        complete_event)
+    expected = [Completion('foo', 0, display_meta='column')]
+    assert set(expected) == set(result)
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index fee7be4d..4063ea7f 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -6,7 +6,7 @@ import pytest
 # Returns the expected select-clause suggestions for a single-table select
 def cols_etc(table, schema=None, alias=None, is_function=False, parent=None):
     return set([
-        Column(tables=((schema, table, alias, is_function),)),
+        Column(table_refs=((schema, table, alias, is_function),)),
         Function(schema=parent),
         Keyword()])
 
@@ -65,13 +65,13 @@ def test_where_equals_any_suggests_columns_or_keywords():
 def test_lparen_suggests_cols():
     suggestion = suggest_type('SELECT MAX( FROM tbl', 'SELECT MAX(')
     assert set(suggestion) == set([
-        Column(tables=((None, 'tbl', None, False),))])
+        Column(table_refs=((None, 'tbl', None, False),))])
 
 
 def test_select_suggests_cols_and_funcs():
     suggestions = suggest_type('SELECT ', 'SELECT ')
     assert set(suggestions) == set([
-         Column(tables=()),
+         Column(table_refs=()),
          Function(schema=None),
          Keyword(),
     ])
@@ -111,7 +111,7 @@ def test_suggest_after_join_with_two_tables(expression):
     suggestions = suggest_type(expression, expression)
     tables = tuple([(None, 'foo', None, False), (None, 'bar', None, False)])
     assert set(suggestions) == set([
-        FromClauseItem(schema=None, tables=tables),
+        FromClauseItem(schema=None, table_refs=tables),
         Join(tables, None),
         Schema(),
     ])
@@ -125,7 +125,7 @@ def test_suggest_after_join_with_one_table(expression):
     suggestions = suggest_type(expression, expression)
     tables = ((None, 'foo', None, False),)
     assert set(suggestions) == set([
-        FromClauseItem(schema=None, tables=tables),
+        FromClauseItem(schema=None, table_refs=tables),
         Join(((None, 'foo', None, False),), None),
         Schema(),
     ])
@@ -174,7 +174,7 @@ def test_suggest_qualified_tables_views_functions_and_joins(expression):
     suggestions = suggest_type(expression, expression)
     tbls = tuple([(None, 'foo', None, False)])
     assert set(suggestions) == set([
-        FromClauseItem(schema='sch', tables=tbls),
+        FromClauseItem(schema='sch', table_refs=tbls),
         Join(tbls, 'sch'),
     ])
 
@@ -197,13 +197,13 @@ def test_truncate_suggests_qualified_tables():
 ])
 def test_distinct_suggests_cols(text):
     suggestions = suggest_type(text, text)
-    assert suggestions ==(Column(tables=()),)
+    assert suggestions ==(Column(table_refs=()),)
 
 
 def test_col_comma_suggests_cols():
     suggestions = suggest_type('SELECT a, b, FROM tbl', 'SELECT a, b,')
     assert set(suggestions) == set([
-        Column(tables=((None, 'tbl', None, False),)),
+        Column(table_refs=((None, 'tbl', None, False),)),
         Function(schema=None),
         Keyword(),
     ])
@@ -232,17 +232,17 @@ def test_into_suggests_tables_and_schemas():
 ])
 def test_insert_into_lparen_suggests_cols(text):
     suggestions = suggest_type(text, 'INSERT INTO abc (')
-    assert suggestions ==(Column(tables=((None, 'abc', None, False),)),)
+    assert suggestions ==(Column(table_refs=((None, 'abc', None, False),)),)
 
 
 def test_insert_into_lparen_partial_text_suggests_cols():
     suggestions = suggest_type('INSERT INTO abc (i', 'INSERT INTO abc (i')
-    assert suggestions ==(Column(tables=((None, 'abc', None, False),)),)
+    assert suggestions ==(Column(table_refs=((None, 'abc', None, False),)),)
 
 
 def test_insert_into_lparen_comma_suggests_cols():
     suggestions = suggest_type('INSERT INTO abc (id,', 'INSERT INTO abc (id,')
-    assert suggestions ==(Column(tables=((None, 'abc', None, False),)),)
+    assert suggestions ==(Column(table_refs=((None, 'abc', None, False),)),)
 
 
 def test_partially_typed_col_name_suggests_col_names():
@@ -254,7 +254,7 @@ def test_partially_typed_col_name_suggests_col_names():
 def test_dot_suggests_cols_of_a_table_or_schema_qualified_table():
     suggestions = suggest_type('SELECT tabl. FROM tabl', 'SELECT tabl.')
     assert set(suggestions) == set([
-        Column(tables=((None, 'tabl', None, False),)),
+        Column(table_refs=((None, 'tabl', None, False),)),
         Table(schema='tabl'),
         View(schema='tabl'),
         Function(schema='tabl'),
@@ -272,7 +272,7 @@ def test_dot_suggests_cols_of_an_alias(sql):
     assert set(suggestions) == set([
         Table(schema='t1'),
         View(schema='t1'),
-        Column(tables=((None, 'tabl1', 't1', False),)),
+        Column(table_refs=((None, 'tabl1', 't1', False),)),
         Function(schema='t1'),
     ])
 
@@ -288,7 +288,7 @@ def test_dot_suggests_cols_of_an_alias_where(sql):
     assert set(suggestions) == set([
         Table(schema='t1'),
         View(schema='t1'),
-        Column(tables=((None, 'tabl1', 't1', False),)),
+        Column(table_refs=((None, 'tabl1', 't1', False),)),
         Function(schema='t1'),
     ])
 
@@ -297,7 +297,7 @@ def test_dot_col_comma_suggests_cols_or_schema_qualified_table():
     suggestions = suggest_type('SELECT t1.a, t2. FROM tabl1 t1, tabl2 t2',
             'SELECT t1.a, t2.')
     assert set(suggestions) == set([
-        Column(tables=((None, 'tabl2', 't2', False),)),
+        Column(table_refs=((None, 'tabl2', 't2', False),)),
         Table(schema='t2'),
         View(schema='t2'),
         Function(schema='t2'),
@@ -328,7 +328,7 @@ def test_outer_table_reference_in_exists_subquery_suggests_columns():
     q = 'SELECT * FROM foo f WHERE EXISTS (SELECT 1 FROM bar WHERE f.'
     suggestions = suggest_type(q, q)
     assert set(suggestions) == set([
-        Column(tables=((None, 'foo', 'f', False),)),
+        Column(table_refs=((None, 'foo', 'f', False),)),
         Table(schema='f'),
         View(schema='f'),
         Function(schema='f'),
@@ -354,7 +354,7 @@ def test_sub_select_table_name_completion_with_outer_table(expression):
     suggestion = suggest_type(expression, expression)
     tbls = tuple([(None, 'foo', None, False)])
     assert set(suggestion) == set([
-        FromClauseItem(schema=None, tables=tbls),
+        FromClauseItem(schema=None, table_refs=tbls),
         Schema(),
     ])
 
@@ -363,7 +363,7 @@ def test_sub_select_col_name_completion():
     suggestions = suggest_type('SELECT * FROM (SELECT  FROM abc',
             'SELECT * FROM (SELECT ')
     assert set(suggestions) == set([
-        Column(tables=((None, 'abc', None, False),)),
+        Column(table_refs=((None, 'abc', None, False),)),
         Function(schema=None),
         Keyword(),
     ])
@@ -380,7 +380,7 @@ def test_sub_select_dot_col_name_completion():
     suggestions = suggest_type('SELECT * FROM (SELECT t. FROM tabl t',
             'SELECT * FROM (SELECT t.')
     assert set(suggestions) == set([
-        Column(tables=((None, 'tabl', 't', False),)),
+        Column(table_refs=((None, 'tabl', 't', False),)),
         Table(schema='t'),
         View(schema='t'),
         Function(schema='t'),
@@ -394,7 +394,7 @@ def test_join_suggests_tables_and_schemas(tbl_alias, join_type):
     suggestion = suggest_type(text, text)
     tbls = tuple([(None, 'abc', tbl_alias or None, False)])
     assert set(suggestion) == set([
-        FromClauseItem(schema=None, tables=tbls),
+        FromClauseItem(schema=None, table_refs=tbls),
         Schema(),
         Join(tbls, None),
     ])
@@ -407,7 +407,7 @@ def test_left_join_with_comma():
     # but there's a bug with commas
     tbls = tuple([(None, 'foo', 'f', False)])
     assert set(suggestions) == set([
-         FromClauseItem(schema=None, tables=tbls),
+         FromClauseItem(schema=None, table_refs=tbls),
          Schema(),
     ])
 
@@ -420,11 +420,11 @@ def test_join_alias_dot_suggests_cols1(sql):
     suggestions = suggest_type(sql, sql)
     tables = ((None, 'abc', 'a', False), (None, 'def', 'd', False))
     assert set(suggestions) == set([
-        Column(tables=((None, 'abc', 'a', False),)),
+        Column(table_refs=((None, 'abc', 'a', False),)),
         Table(schema='a'),
         View(schema='a'),
         Function(schema='a'),
-        JoinCondition(tables=tables, parent=(None, 'abc', 'a', False))
+        JoinCondition(table_refs=tables, parent=(None, 'abc', 'a', False))
     ])
 
 
@@ -435,7 +435,7 @@ def test_join_alias_dot_suggests_cols1(sql):
 def test_join_alias_dot_suggests_cols2(sql):
     suggestion = suggest_type(sql, sql)
     assert set(suggestion) == set([
-        Column(tables=((None, 'def', 'd', False),)),
+        Column(table_refs=((None, 'def', 'd', False),)),
         Table(schema='d'),
         View(schema='d'),
         Function(schema='d'),
@@ -457,7 +457,7 @@ on ''',
 def test_on_suggests_aliases_and_join_conditions(sql):
     suggestions = suggest_type(sql, sql)
     tables = ((None, 'abc', 'a', False), (None, 'bcd', 'b', False))
-    assert set(suggestions) == set((JoinCondition(tables=tables, parent=None),
+    assert set(suggestions) == set((JoinCondition(table_refs=tables, parent=None),
         Alias(aliases=('a', 'b',)),))
 
 @pytest.mark.parametrize('sql', [
@@ -467,7 +467,7 @@ def test_on_suggests_aliases_and_join_conditions(sql):
 def test_on_suggests_tables_and_join_conditions(sql):
     suggestions = suggest_type(sql, sql)
     tables = ((None, 'abc', None, False), (None, 'bcd', None, False))
-    assert set(suggestions) == set((JoinCondition(tables=tables, parent=None),
+    assert set(suggestions) == set((JoinCondition(table_refs=tables, parent=None),
         Alias(aliases=('abc', 'bcd',)),))
 
 
@@ -487,7 +487,7 @@ def test_on_suggests_aliases_right_side(sql):
 def test_on_suggests_tables_and_join_conditions_right_side(sql):
     suggestions = suggest_type(sql, sql)
     tables = ((None, 'abc', None, False), (None, 'bcd', None, False))
-    assert set(suggestions) == set((JoinCondition(tables=tables, parent=None),
+    assert set(suggestions) == set((JoinCondition(table_refs=tables, parent=None),
       Alias(aliases=('abc', 'bcd',)),))
 
 
@@ -503,7 +503,7 @@ def test_on_suggests_tables_and_join_conditions_right_side(sql):
 def test_join_using_suggests_common_columns(text):
     tables = ((None, 'abc', None, False), (None, 'def', None, False))
     assert set(suggest_type(text, text)) == set([
-        Column(tables=tables, require_last_table=True),])
+        Column(table_refs=tables, require_last_table=True),])
 
 
 def test_suggest_columns_after_multiple_joins():
@@ -513,7 +513,7 @@ def test_suggest_columns_after_multiple_joins():
             inner join t3 ON
               t2.id = t3.'''
     suggestions = suggest_type(sql, sql)
-    assert Column(tables=((None, 't3', None, False),)) in set(suggestions)
+    assert Column(table_refs=((None, 't3', None, False),)) in set(suggestions)
 
 
 def test_2_statements_2nd_current():
@@ -527,7 +527,7 @@ def test_2_statements_2nd_current():
     suggestions = suggest_type('select * from a; select  from b',
                                'select * from a; select ')
     assert set(suggestions) == set([
-        Column(tables=((None, 'b', None, False),)),
+        Column(table_refs=((None, 'b', None, False),)),
         Function(schema=None),
         Keyword()
     ])
@@ -684,7 +684,7 @@ def test_suggest_where_keyword(text):
 
 @pytest.mark.parametrize('text, before, expected', [
     ('\\ns abc SELECT ', 'SELECT ', [
-        Column(tables=()),
+        Column(table_refs=()),
         Function(schema=None),
         Keyword()
     ]),
@@ -692,7 +692,7 @@ def test_suggest_where_keyword(text):
     ('\\ns abc SELECT t1. FROM tabl1 t1', 'SELECT t1.', [
         Table(schema='t1'),
         View(schema='t1'),
-        Column(tables=((None, 'tabl1', 't1', False),)),
+        Column(table_refs=((None, 'tabl1', 't1', False),)),
         Function(schema='t1')
     ])
 ])
@@ -725,6 +725,6 @@ def test_ignore_leading_double_quotes(sql):
 def test_column_keyword_suggests_columns(sql):
     suggestions = suggest_type(sql, sql)
     assert set(suggestions) == set([
-        Column(tables=((None, 'foo', None, False),)),
+        Column(table_refs=((None, 'foo', None, False),)),
     ])