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:
parent
ad2eb9f976
commit
e060dd398e
@ -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':
|
||||
|
@ -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': []},
|
||||
|
Loading…
Reference in New Issue
Block a user