def test_select_with_hanging_comma_single_table():
tables = extract_tables('select a, from abc')
assert tables == ((None, 'abc', None, False), )
def test_select_with_hanging_period_multiple_tables():
tables = extract_tables('SELECT t1. FROM tabl1 t1, tabl2 t2')
assert set(tables) == set([(None, 'tabl1', 't1', False),
(None, 'tabl2', 't2', False)])
def test_simple_select_with_cols_single_table():
tables = extract_tables('select a,b from abc')
assert tables == ((None, 'abc', None, False), )
def test_simple_select_with_cols_multiple_tables():
tables = extract_tables('select a,b from abc, def')
assert set(tables) == set([(None, 'abc', None, False),
(None, 'def', None, False)])
def test_empty_string():
tables = extract_tables('')
assert tables == ()
def test_simple_select_multiple_tables_deouble_quoted_aliased():
tables = extract_tables('select * from "Abc" a, "Def" d')
assert set(tables) == set([(None, 'Abc', 'a', False),
(None, 'Def', 'd', False)])
def test_incomplete_join_clause():
sql = '''select a.x, b.y
from abc a join bcd b
on a.id = '''
tables = extract_tables(sql)
assert tables == ((None, 'abc', 'a', False), (None, 'bcd', 'b', False))
def test_simple_select_single_table_double_quoted():
tables = extract_tables('select * from "Abc"')
assert tables == ((None, 'Abc', None, False), )
def test_join_table(join_type):
sql = 'SELECT * FROM abc a {0} JOIN def d ON a.id = d.num'.format(
join_type)
tables = extract_tables(sql)
assert set(tables) == set([(None, 'abc', 'a', False),
(None, 'def', 'd', False)])
def test_join_table_schema_qualified():
tables = extract_tables(
'SELECT * FROM abc.def x JOIN ghi.jkl y ON x.id = y.num')
assert set(tables) == set([('abc', 'def', 'x', False),
('ghi', 'jkl', 'y', False)])
def test_simple_update_table_with_schema():
tables = extract_tables('update abc.def set id = 1')
assert tables == (('abc', 'def', None, False), )
def test_simple_update_table_no_schema():
tables = extract_tables('update abc set id = 1')
assert tables == ((None, 'abc', None, False), )
def test_simple_insert_single_table_schema_qualified():
tables = extract_tables('insert into abc.def (id, name) values (1, "def")')
assert tables == (('abc', 'def', None, False), )
def test_complex_table_and_function():
tables = extract_tables('''SELECT * FROM foo.bar baz
JOIN bar.qux(x, y, z) quux''')
assert set(tables) == set([('foo', 'bar', 'baz', False),
('bar', 'qux', 'quux', True)])
def test_join_as_table():
tables = extract_tables('SELECT * FROM my_table AS m WHERE m.a > 5')
assert tables == ((None, 'my_table', 'm', False), )
def test_simple_select_single_table_schema_qualified(sql):
tables = extract_tables(sql)
assert tables == (('abc', 'def', None, False), )
def test_subselect_tables():
sql = 'SELECT * FROM (SELECT FROM abc'
tables = extract_tables(sql)
assert tables == ((None, 'abc', None, False), )
def test_simple_select_multiple_tables_double_quoted():
tables = extract_tables('select * from "Abc", "Def"')
assert set(tables) == set([(None, 'Abc', None, False),
(None, 'Def', None, False)])
def test_extract_no_tables(text):
tables = extract_tables(text)
assert tables == tuple()
def test_simple_select_single_table_deouble_quoted_aliased():
tables = extract_tables('select * from "Abc" a')
assert tables == ((None, 'Abc', 'a', False), )
def test_simple_select_single_table_schema_qualified_quoted_table(sql):
tables = extract_tables(sql)
assert tables == (('abc', 'def', '"def"', False), )
def test_simple_select_multiple_tables_schema_qualified():
tables = extract_tables('select * from abc.def, ghi.jkl')
assert set(tables) == set([('abc', 'def', None, False),
('ghi', 'jkl', None, False)])
def test_simple_schema_qualified_function_as_table(arg_list):
tables = extract_tables('SELECT * FROM foo.bar({0})'.format(arg_list))
assert tables == (('foo', 'bar', None, True), )
def test_simple_select_with_cols_single_table_schema_qualified():
tables = extract_tables('select a,b from abc.def')
assert tables == (('abc', 'def', None, False), )
def test_simple_aliased_function_as_table(arg_list):
tables = extract_tables('SELECT * FROM foo({0}) bar'.format(arg_list))
assert tables == ((None, 'foo', 'bar', True), )
def test_simple_select_with_cols_multiple_qualified_tables():
tables = extract_tables('select a,b from abc.def, def.ghi')
assert set(tables) == set([('abc', 'def', None, False),
('def', 'ghi', None, False)])
def test_simple_table_and_function():
tables = extract_tables('SELECT * FROM foo JOIN bar()')
assert set(tables) == set([(None, 'foo', None, False),
(None, 'bar', None, True)])
def test_select_with_hanging_comma_multiple_tables():
tables = extract_tables('select a, from abc, def')
assert set(tables) == set([(None, 'abc', None, False),
(None, 'def', None, False)])
def suggest_based_on_last_token(token, stmt):
if isinstance(token, string_types):
token_v = token.lower()
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.
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 = stmt.reduce_to_prev_keyword()
return suggest_based_on_last_token(prev_keyword, stmt)
elif isinstance(token, Identifier):
# If the previous token is an identifier, we can suggest datatypes if
# we're in a parenthesized column/field list, e.g.:
# CREATE TABLE foo (Identifier <CURSOR>
# CREATE FUNCTION foo (Identifier <CURSOR>
# If we're not in a parenthesized list, the most likely scenario is the
# user is about to specify an alias, e.g.:
# SELECT Identifier <CURSOR>
# SELECT foo FROM Identifier <CURSOR>
prev_keyword, _ = find_prev_keyword(stmt.text_before_cursor)
if prev_keyword and prev_keyword.value == '(':
# Suggest datatypes
return suggest_based_on_last_token('type', stmt)
return (Keyword(), )
else:
token_v = token.value.lower()
if not token:
return (Keyword(), Special())
if token_v.endswith('('):
p = sqlparse.parse(stmt.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', stmt)
# Check for a subquery expression (cases 3 & 4)
where = p.tokens[-1]
prev_tok = where.token_prev(len(where.tokens) - 1)[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 (Keyword(), )
return column_suggestions
# Get the token before the parens
prev_tok = p.token_prev(len(p.tokens) - 1)[1]
if (prev_tok and prev_tok.value
and prev_tok.value.lower().split(' ')[-1] == 'using'):
# tbl1 INNER JOIN tbl2 USING (col1, col2)
tables = stmt.get_tables('before')
# suggest columns that are present in more than one table
return (Column(table_refs=tables,
require_last_table=True,
local_tables=stmt.local_tables), )
if p.token_first().value.lower() == 'select':
# If the lparen is preceeded by a space chances are we're about to
# do a sub-select.
if last_word(stmt.text_before_cursor,
'all_punctuations').startswith('('):
return (Keyword(), )
prev_prev_tok = prev_tok and p.token_prev(p.token_index(prev_tok))[1]
if prev_prev_tok and prev_prev_tok.normalized == 'INTO':
return (Column(table_refs=stmt.get_tables('insert'),
context='insert'), )
# We're probably in a function argument list
return (Column(table_refs=extract_tables(stmt.full_text),
local_tables=stmt.local_tables,
qualifiable=True), )
if token_v == 'set':
return (Column(table_refs=stmt.get_tables(),
local_tables=stmt.local_tables), )
if token_v in ('select', 'where', 'having', 'by', 'distinct'):
# Check for a table alias or schema qualification
parent = stmt.identifier.get_parent_name() \
if (stmt.identifier and stmt.identifier.get_parent_name()) else []
tables = stmt.get_tables()
if parent:
tables = tuple(t for t in tables if identifies(parent, t))
return (
Column(table_refs=tables, local_tables=stmt.local_tables),
Table(schema=parent),
View(schema=parent),
Function(schema=parent),
)
return (
Column(table_refs=tables,
local_tables=stmt.local_tables,
qualifiable=True),
Function(schema=None),
Keyword(token_v.upper()),
)
if token_v == 'as':
# Don't suggest anything for aliases
return ()
if (token_v.endswith('join') and token.is_keyword) or \
token_v in ('copy', 'from', 'update', 'into', 'describe', 'truncate'):
schema = stmt.get_identifier_schema()
tables = extract_tables(stmt.text_before_cursor)
is_join = token_v.endswith('join') and token.is_keyword
# Suggest tables from either the currently-selected schema or the
# public schema if no schema has been specified
suggest = []
if not schema:
# Suggest schemas
suggest.insert(0, Schema())
if token_v == 'from' or is_join:
suggest.append(
FromClauseItem(schema=schema,
table_refs=tables,
local_tables=stmt.local_tables))
elif token_v == 'truncate':
suggest.append(Table(schema))
else:
suggest.extend((Table(schema), View(schema)))
if is_join and _allow_join(stmt.parsed):
tables = stmt.get_tables('before')
suggest.append(Join(table_refs=tables, schema=schema))
return tuple(suggest)
if token_v == 'function':
schema = stmt.get_identifier_schema()
# stmt.get_previous_token will fail for e.g. `SELECT 1 FROM functions
# WHERE function:`
try:
prev = stmt.get_previous_token(token).value.lower()
if prev in ('drop', 'alter', 'create', 'create or replace'):
return (Function(schema=schema, usage='signature'), )
except ValueError:
pass
return tuple()
if token_v in ('table', 'view'):
# E.g. 'ALTER TABLE <tablname>'
rel_type = {
'table': Table,
'view': View,
'function': Function
}[token_v]
schema = stmt.get_identifier_schema()
if schema:
return (rel_type(schema=schema), )
return (Schema(), rel_type(schema=schema))
if token_v == 'column':
# E.g. 'ALTER TABLE foo ALTER COLUMN bar
return (Column(table_refs=stmt.get_tables()), )
if token_v == 'on':
tables = stmt.get_tables('before')
parent = stmt.identifier.get_parent_name() \
if (stmt.identifier and stmt.identifier.get_parent_name()) else None
if parent:
# "ON parent.<suggestion>"
# parent can be either a schema name or table alias
filteredtables = tuple(t for t in tables if identifies(parent, t))
sugs = [
Column(table_refs=filteredtables,
local_tables=stmt.local_tables),
Table(schema=parent),
View(schema=parent),
Function(schema=parent)
]
if filteredtables and _allow_join_condition(stmt.parsed):
sugs.append(
JoinCondition(table_refs=tables,
parent=filteredtables[-1]))
return tuple(sugs)
# ON <suggestion>
# Use table alias if there is one, otherwise the table name
aliases = tuple(t.ref for t in tables)
if _allow_join_condition(stmt.parsed):
return (Alias(aliases=aliases),
JoinCondition(table_refs=tables, parent=None))
return (Alias(aliases=aliases), )
if token_v in ('c', 'use', 'database', 'template'):
# "\c <db", "use <db>", "DROP DATABASE <db>",
# "CREATE DATABASE <newdb> WITH TEMPLATE <db>"
return (Database(), )
if token_v == 'schema':
# DROP SCHEMA schema_name, SET SCHEMA schema name
prev_keyword = stmt.reduce_to_prev_keyword(n_skip=2)
quoted = prev_keyword and prev_keyword.value.lower() == 'set'
return (Schema(quoted), )
if token_v.endswith(',') or token_v in ('=', 'and', 'or'):
prev_keyword = stmt.reduce_to_prev_keyword()
if prev_keyword:
return suggest_based_on_last_token(prev_keyword, stmt)
return ()
if token_v in ('type', '::'):
# ALTER TABLE foo SET DATA TYPE bar
# SELECT foo::bar
# Note that tables are a form of composite type in postgresql, so
# they're suggested here as well
schema = stmt.get_identifier_schema()
suggestions = [Datatype(schema=schema), Table(schema=schema)]
if not schema:
suggestions.append(Schema())
return tuple(suggestions)
if token_v in {'alter', 'create', 'drop'}:
return (Keyword(token_v.upper()), )
if token.is_keyword:
# token is a keyword we haven't implemented any special handling for
# go backwards in the query until we find one we do recognize
prev_keyword = stmt.reduce_to_prev_keyword(n_skip=1)
if prev_keyword:
return suggest_based_on_last_token(prev_keyword, stmt)
return (Keyword(token_v.upper()), )
return (Keyword(), )
