Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
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_multiple_joins():
sql = """select * from t1
inner join t2 ON
t1.id = t2.t1_id
inner join t3 ON
t2.id = t3."""
tables = extract_tables(sql)
assert tables == (
(None, "t1", None, False),
(None, "t2", None, False),
(None, "t3", None, False),
)
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_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_select_with_hanging_comma_single_table():
tables = extract_tables("select a, from abc")
assert tables == ((None, "abc", None, False),)
def test_simple_select_single_table_double_quoted():
tables = extract_tables('select * from "Abc"')
assert tables == ((None, 'Abc', None, False),)
def test_simple_select_multiple_tables():
tables = extract_tables('select * from abc, def')
assert set(tables) == set([(None, 'abc', None, False),
(None, '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)]
)
local_tables=stmt.local_tables,
),
)
elif 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,
),
)
elif token_v == "set":
return (Column(table_refs=stmt.get_tables(), local_tables=stmt.local_tables),)
elif token_v in ("select", "where", "having", "order by", "distinct"):
# Check for a table alias or schema qualification
parent = (stmt.identifier and stmt.identifier.get_parent_name()) or []
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),
return (
Column(
table_refs=tables, local_tables=stmt.local_tables, qualifiable=True
),
Function(schema=None),
Keyword(token_v.upper()),
)
elif token_v == "as":
# Don't suggest anything for aliases
return ()
elif (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
)
)