1
1
mirror of https://github.com/dbcli/pgcli.git synced 2024-10-06 02:07:53 +03:00

improve suggestions in non-trivial WHERE clauses

sqlparse groups tokens belonging to the where clause into a single Where token (a subclass of TokenList). In order to handle cases beyond the simplest `SELECT * FROM foo WHERE`, we need to look "inside" of this token list.
This commit is contained in:
Darik Gamble 2015-05-02 09:43:41 -04:00
parent ad2eb9f976
commit e060dd398e
2 changed files with 115 additions and 16 deletions

View File

@ -1,7 +1,7 @@
from __future__ import print_function
import sys
import sqlparse
from sqlparse.sql import Comparison, Identifier
from sqlparse.sql import Comparison, Identifier, Where
from .parseutils import last_word, extract_tables, find_prev_keyword
from .pgspecial import parse_special_command
@ -134,23 +134,69 @@ def suggest_special(text):
def suggest_based_on_last_token(token, text_before_cursor, full_text, identifier):
if isinstance(token, string_types):
token_v = token.lower()
else:
elif isinstance(token, Comparison):
# 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.lower()
else:
token_v = token.value.lower()
token_v = token.tokens[-1].value.lower()
elif isinstance(token, Where):
# sqlparse groups all tokens from the where clause into a single token
# list. This means that token.value may be something like
# 'where foo > 5 and '. We need to look "inside" token.tokens to handle
# suggestions in complicated where clauses correctly
prev_keyword, text_before_cursor = find_prev_keyword(text_before_cursor)
return suggest_based_on_last_token(prev_keyword, text_before_cursor,
full_text, identifier)
else:
token_v = token.value.lower()
if not token:
return [{'type': 'keyword'}, {'type': 'special'}]
elif token_v.endswith('('):
p = sqlparse.parse(text_before_cursor)[0]
if p.tokens and isinstance(p.tokens[-1], Where):
# Four possibilities:
# 1 - Parenthesized clause like "WHERE foo AND ("
# Suggest columns/functions
# 2 - Function call like "WHERE foo("
# Suggest columns/functions
# 3 - Subquery expression like "WHERE EXISTS ("
# Suggest keywords, in order to do a subquery
# 4 - Subquery OR array comparison like "WHERE foo = ANY("
# Suggest columns/functions AND keywords. (If we wanted to be
# really fancy, we could suggest only array-typed columns)
column_suggestions = suggest_based_on_last_token('where',
text_before_cursor, full_text, identifier)
# Check for a subquery expression (cases 3 & 4)
where = p.tokens[-1]
prev_tok = where.token_prev(len(where.tokens) - 1)
if isinstance(prev_tok, Comparison):
# e.g. "SELECT foo FROM bar WHERE foo = ANY("
prev_tok = prev_tok.tokens[-1]
prev_tok = prev_tok.value.lower()
if prev_tok == 'exists':
return [{'type': 'keyword'}]
elif prev_tok in ('any', 'some', 'all'):
return column_suggestions + [{'type': 'keyword'}]
elif prev_tok == 'in':
# Technically, we should suggest columns AND keywords, as
# per case 4. However, IN is different from ANY, SOME, ALL
# in that it can accept a *list* of columns, or a subquery.
# Because of the way lists of columns are reduced in this
# function, "SELECT * FROM foo WHERE bar IN (baz, qux, " would
# also suggest keywords, which is kinda gross.
return column_suggestions
else:
return column_suggestions
# Get the token before the parens
prev_tok = p.token_prev(len(p.tokens) - 1)
if prev_tok and prev_tok.value and prev_tok.value.lower() == 'using':

View File

@ -17,13 +17,43 @@ def test_select_suggests_cols_with_qualified_table_scope():
{'type': 'column', 'tables': [('sch', 'tabl', None)]},
{'type': 'function', 'schema': []}])
def test_where_suggests_columns_functions():
suggestions = suggest_type('SELECT * FROM tabl WHERE ',
'SELECT * FROM tabl WHERE ')
@pytest.mark.parametrize('expression', [
'SELECT * FROM tabl WHERE ',
'SELECT * FROM tabl WHERE (',
'SELECT * FROM tabl WHERE foo = ',
'SELECT * FROM tabl WHERE bar OR ',
'SELECT * FROM tabl WHERE foo = 1 AND ',
'SELECT * FROM tabl WHERE (bar > 10 AND ',
'SELECT * FROM tabl WHERE (bar AND (baz OR (qux AND (',
'SELECT * FROM tabl WHERE 10 < ',
'SELECT * FROM tabl WHERE foo BETWEEN ',
'SELECT * FROM tabl WHERE foo BETWEEN foo AND ',
])
def test_where_suggests_columns_functions(expression):
suggestions = suggest_type(expression, expression)
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'column', 'tables': [(None, 'tabl', None)]},
{'type': 'function', 'schema': []}])
@pytest.mark.parametrize('expression', [
'SELECT * FROM tabl WHERE foo IN (',
'SELECT * FROM tabl WHERE foo IN (bar, ',
])
def test_where_in_suggests_columns(expression):
suggestions = suggest_type(expression, expression)
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'column', 'tables': [(None, 'tabl', None)]},
{'type': 'function', 'schema': []}])
def test_where_equals_any_suggests_columns_or_keywords():
text = 'SELECT * FROM tabl WHERE foo = ANY('
suggestions = suggest_type(text, text)
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'column', 'tables': [(None, 'tabl', None)]},
{'type': 'function', 'schema': []},
{'type': 'keyword'}])
def test_lparen_suggests_cols():
suggestion = suggest_type('SELECT MAX( FROM tbl', 'SELECT MAX(')
assert suggestion == [
@ -145,17 +175,40 @@ def test_dot_col_comma_suggests_cols_or_schema_qualified_table():
{'type': 'view', 'schema': 't2'},
{'type': 'function', 'schema': 't2'}])
def test_sub_select_suggests_keyword():
suggestion = suggest_type('SELECT * FROM (', 'SELECT * FROM (')
@pytest.mark.parametrize('expression', [
'SELECT * FROM (',
'SELECT * FROM foo WHERE EXISTS (',
'SELECT * FROM foo WHERE bar AND NOT EXISTS (',
])
def test_sub_select_suggests_keyword(expression):
suggestion = suggest_type(expression, expression)
assert suggestion == [{'type': 'keyword'}]
def test_sub_select_partial_text_suggests_keyword():
suggestion = suggest_type('SELECT * FROM (S', 'SELECT * FROM (S')
@pytest.mark.parametrize('expression', [
'SELECT * FROM (S',
'SELECT * FROM foo WHERE EXISTS (S',
'SELECT * FROM foo WHERE bar AND NOT EXISTS (S',
])
def test_sub_select_partial_text_suggests_keyword(expression):
suggestion = suggest_type(expression, expression)
assert suggestion == [{'type': 'keyword'}]
def test_sub_select_table_name_completion():
suggestion = suggest_type('SELECT * FROM (SELECT * FROM ',
'SELECT * FROM (SELECT * FROM ')
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 suggestions == [
{'type': 'column', 'tables': [(None, 'foo', 'f')]},
{'type': 'table', 'schema': 'f'},
{'type': 'view', 'schema': 'f'},
{'type': 'function', 'schema': 'f'}]
@pytest.mark.parametrize('expression', [
'SELECT * FROM (SELECT * FROM ',
'SELECT * FROM foo WHERE EXISTS (SELECT * FROM ',
'SELECT * FROM foo WHERE bar AND NOT EXISTS (SELECT * FROM ',
])
def test_sub_select_table_name_completion(expression):
suggestion = suggest_type(expression, expression)
assert sorted_dicts(suggestion) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},