def adapter(data, headers, table_format=None, **kwargs):
tables = extract_tables(formatter.query)
if len(tables) > 0:
table = tables[0]
if table[0]:
table_name = "{}.{}".format(*table[:2])
else:
table_name = table[1]
else:
table_name = "`DUAL`"
if table_format == 'sql-insert':
h = "`, `".join(headers)
yield "INSERT INTO {} (`{}`) VALUES".format(table_name, h)
prefix = " "
for d in data:
values = ", ".join(
escape_for_sql_statement(v) for i, v in enumerate(d))
yield "{}({})".format(prefix, values)
if prefix == " ":
prefix = ", "
yield ";"
if table_format.startswith('sql-update'):
s = table_format.split('-')
keys = 1
if len(s) > 2:
keys = int(s[-1])
for d in data:
yield "UPDATE {} SET".format(table_name)
prefix = " "
for i, v in enumerate(d[keys:], keys):
yield "{}`{}` = {}".format(prefix, headers[i],
escape_for_sql_statement(v))
if prefix == " ":
prefix = ", "
f = "`{}` = {}"
where = (f.format(headers[i], escape_for_sql_statement(d[i]))
for i in range(keys))
yield "WHERE {};".format(" AND ".join(where))
def adapter(data, headers, table_format=None, **kwargs):
tables = extract_tables(formatter.query)
if len(tables) > 0:
table = tables[0]
if table[0]:
table_name = "{}.{}".format(*table[:2])
else:
table_name = table[1]
else:
table_name = "`DUAL`"
escape = formatter.mycli.sqlexecute.conn.escape
if table_format == 'sql-insert':
h = "`, `".join(headers)
yield "INSERT INTO {} (`{}`) VALUES".format(table_name, h)
prefix = " "
for d in data:
values = ", ".join(escape(v) for i, v in enumerate(d))
yield "{}({})".format(prefix, values)
if prefix == " ":
prefix = ", "
yield ";"
if table_format.startswith('sql-update'):
s = table_format.split('-')
keys = 1
if len(s) > 2:
keys = int(s[-1])
for d in data:
yield "UPDATE {} SET".format(table_name)
prefix = " "
for i, v in enumerate(d[keys:], keys):
yield "{}`{}` = {}".format(prefix, headers[i], escape(v))
if prefix == " ":
prefix = ", "
f = "`{}` = {}"
where = (f.format(headers[i], escape(d[i])) for i in range(keys))
yield "WHERE {};".format(" AND ".join(where))
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)]
def test_simple_insert_single_table():
tables = extract_tables('insert into abc (id, name) values (1, "def")')
# sqlparse mistakenly assigns an alias to the table
# assert tables == [(None, 'abc', None)]
assert tables == [(None, 'abc', 'abc')]
def test_select_with_hanging_comma_multiple_tables():
tables = extract_tables('select a, from abc, def')
assert sorted(tables) == [(None, 'abc', None), (None, 'def', None)]
def test_simple_select_with_cols_multiple_tables_with_schema():
tables = extract_tables('select a,b from abc.def, def.ghi')
assert sorted(tables) == [('abc', 'def', None), ('def', 'ghi', None)]
def test_simple_select_with_cols_single_table_schema_qualified():
tables = extract_tables('select a,b from abc.def')
assert tables == [('abc', 'def', None)]
def test_simple_select_multiple_tables_schema_qualified():
tables = extract_tables('select * from abc.def, ghi.jkl')
assert sorted(tables) == [('abc', 'def', None), ('ghi', 'jkl', None)]
def __extract_tables(self):
"""获取sql语句中的表名"""
return [
list(i)[1].replace('`', '').lower()
for i in extract_tables(self.sqltext)
]
def test_select_with_hanging_comma_single_table():
tables = extract_tables('select a, from abc')
assert tables == [(None, 'abc', None)]
def test_empty_string():
tables = extract_tables('')
assert tables == []
def test_simple_select_with_cols_multiple_tables_with_schema():
tables = extract_tables('select a,b from abc.def, def.ghi')
assert sorted(tables) == [('abc', 'def', None), ('def', 'ghi', None)]
def test_simple_select_with_cols_multiple_tables():
tables = extract_tables('select a,b from abc, def')
assert sorted(tables) == [(None, 'abc', None), (None, 'def', None)]
def test_simple_select_with_cols_single_table_schema_qualified():
tables = extract_tables('select a,b from abc.def')
assert tables == [('abc', 'def', None)]
def test_simple_select_multiple_tables_schema_qualified():
tables = extract_tables('select * from abc.def, ghi.jkl')
assert sorted(tables) == [('abc', 'def', None), ('ghi', 'jkl', None)]
def test_simple_update_table_with_schema():
tables = extract_tables('update abc.def set id = 1')
assert tables == [('abc', 'def', None)]
def test_join_table_schema_qualified():
tables = extract_tables(
'SELECT * FROM abc.def x JOIN ghi.jkl y ON x.id = y.num')
assert tables == [('abc', 'def', 'x'), ('ghi', 'jkl', 'y')]
def test_select_with_hanging_comma_multiple_tables():
tables = extract_tables('select a, from abc, def')
assert sorted(tables) == [(None, 'abc', None), (None, 'def', None)]
def test_join_as_table():
tables = extract_tables('SELECT * FROM my_table AS m WHERE m.a > 5')
assert tables == [(None, 'my_table', 'm')]
def test_select_with_hanging_period_multiple_tables():
tables = extract_tables('SELECT t1. FROM tabl1 t1, tabl2 t2')
assert sorted(tables) == [(None, 'tabl1', 't1'), (None, 'tabl2', 't2')]
def test_simple_select_with_cols_single_table():
tables = extract_tables('select a,b from abc')
assert tables == [(None, 'abc', None)]
def test_simple_insert_single_table():
tables = extract_tables('insert into abc (id, name) values (1, "def")')
# sqlparse mistakenly assigns an alias to the table
# assert tables == [(None, 'abc', None)]
assert tables == [(None, 'abc', 'abc')]
def test_simple_select_with_cols_multiple_tables():
tables = extract_tables('select a,b from abc, def')
assert sorted(tables) == [(None, 'abc', None), (None, 'def', None)]
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)]
def test_select_with_hanging_comma_single_table():
tables = extract_tables('select a, from abc')
assert tables == [(None, 'abc', None)]
def test_simple_update_table():
tables = extract_tables('update abc set id = 1')
assert tables == [(None, 'abc', None)]
def test_select_with_hanging_period_multiple_tables():
tables = extract_tables('SELECT t1. FROM tabl1 t1, tabl2 t2')
assert sorted(tables) == [(None, 'tabl1', 't1'), (None, 'tabl2', 't2')]
def test_simple_update_table_with_schema():
tables = extract_tables('update abc.def set id = 1')
assert tables == [('abc', 'def', None)]
def test_empty_string():
tables = extract_tables('')
assert tables == []
def test_join_table():
tables = extract_tables('SELECT * FROM abc a JOIN def d ON a.id = d.num')
assert sorted(tables) == [(None, 'abc', 'a'), (None, 'def', 'd')]
def test_simple_update_table():
tables = extract_tables('update abc set id = 1')
assert tables == [(None, 'abc', None)]
def test_join_table_schema_qualified():
tables = extract_tables('SELECT * FROM abc.def x JOIN ghi.jkl y ON x.id = y.num')
assert tables == [('abc', 'def', 'x'), ('ghi', 'jkl', 'y')]
def test_join_table():
tables = extract_tables('SELECT * FROM abc a JOIN def d ON a.id = d.num')
assert sorted(tables) == [(None, 'abc', 'a'), (None, 'def', 'd')]
def test_simple_select_single_table():
tables = extract_tables('select * from abc')
assert tables == [(None, 'abc', None)]
def test_join_as_table():
tables = extract_tables('SELECT * FROM my_table AS m WHERE m.a > 5')
assert tables == [(None, 'my_table', 'm')]
def suggest_based_on_last_token(token, text_before_cursor, full_text, identifier):
if isinstance(token, str):
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, 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()
is_operand = lambda x: x and any([x.endswith(op) for op in ["+", "-", "*", "/"]])
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]
idx, 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"}]
else:
return column_suggestions
# Get the token before the parens
idx, prev_tok = p.token_prev(len(p.tokens) - 1)
if prev_tok and prev_tok.value and prev_tok.value.lower() == "using":
# tbl1 INNER JOIN tbl2 USING (col1, col2)
tables = extract_tables(full_text)
# suggest columns that are present in more than one table
return [{"type": "column", "tables": tables, "drop_unique": True}]
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(text_before_cursor, "all_punctuations").startswith("("):
return [{"type": "keyword"}]
elif p.token_first().value.lower() == "show":
return [{"type": "show"}]
first_token = p.token_first()
for t in p.tokens[::-1]:
if (t.ttype == Keyword and t.value.lower() == "values") or (
t.is_group and t.tokens[0].value.lower() == "values"
):
# for statement like 「insert into table_name values ( 」 would suggest nothing
if (
first_token
and first_token.ttype == DML
and first_token.value.lower() == "insert"
):
return [{"type": "column_hint"}]
# We're probably in a function argument list
return [{"type": "column", "tables": extract_tables(full_text)}]
elif token_v in ("set", "order by", "distinct"):
return [{"type": "column", "tables": extract_tables(full_text)}]
elif token_v == "as":
# Don't suggest anything for an alias
return []
elif token_v in ("show"):
return [{"type": "show"}]
elif token_v in ("to",):
p = sqlparse.parse(text_before_cursor)[0]
if p.token_first().value.lower() == "change":
return [{"type": "change"}]
else:
return [{"type": "user"}]
elif token_v in ("user", "for"):
# for edge cases get_suggestions("select user fro", len("select user fro"))
if isinstance(token, Token):
if token.parent.token_first().value.lower() == "select":
token_v = "select"
return suggest_based_on_last_token(
"select", text_before_cursor, full_text, identifier
)
return [{"type": "user"}]
elif token_v in ("select", "where", "having"):
# Check for a table alias or schema qualification
parent = (identifier and identifier.get_parent_name()) or []
tables = extract_tables(full_text)
if parent:
tables = [t for t in tables if identifies(parent, *t)]
return [
{"type": "column", "tables": tables},
{"type": "table", "schema": parent},
{"type": "view", "schema": parent},
{"type": "function", "schema": parent},
]
rv = [
{"type": "column", "tables": tables},
{"type": "function", "schema": []},
{"type": "keyword"},
]
if token_v == "select":
return rv
aliases = [alias or table for (schema, table, alias) in tables]
rv.append({"type": "alias", "aliases": aliases})
return rv
elif (token_v.endswith("join") and token.is_keyword) or (
token_v in suggestion_table_tuple
):
schema = (identifier and identifier.get_parent_name()) or []
# Suggest tables from either the currently-selected schema or the
# public schema if no schema has been specified
suggest = [{"type": "table", "schema": schema}]
if not schema:
# Suggest schemas
suggest.insert(0, {"type": "schema"})
# Only tables can be TRUNCATED, otherwise suggest views
if token_v != "truncate":
suggest.append({"type": "view", "schema": schema})
# suggest database when not appear database in select, insert into ... like statement
if len(schema) == 0:
suggest.append({"type": "database", "table": ""})
return suggest
elif token_v in ("table", "view", "function"):
# E.g. 'DROP FUNCTION <funcname>', 'ALTER TABLE <tablname>'
rel_type = token_v
schema = (identifier and identifier.get_parent_name()) or []
if schema:
return [{"type": rel_type, "schema": schema}]
else:
return [{"type": rel_type, "schema": []}, {"type": "database"}]
elif token_v == "on":
tables = extract_tables(full_text) # [(schema, table, alias), ...]
parent = (identifier and identifier.get_parent_name()) or []
if parent:
# "ON parent.<suggestion>"
# parent can be either a schema name or table alias
tables = [t for t in tables if identifies(parent, *t)]
return [
{"type": "column", "tables": tables},
{"type": "table", "schema": parent},
{"type": "view", "schema": parent},
{"type": "function", "schema": parent},
]
else:
# ON <suggestion>
# Use table alias if there is one, otherwise the table name
aliases = [alias or table for (schema, table, alias) in tables]
suggest = [{"type": "alias", "aliases": aliases}]
# The lists of 'aliases' could be empty if we're trying to complete
# a GRANT query. eg: GRANT SELECT, INSERT ON <tab>
# In that case we just suggest all tables.
if not aliases:
suggest.append({"type": "table", "schema": parent})
return suggest
elif token_v in ("use", "database", "template", "connect"):
# "\c <db", "use <db>", "DROP DATABASE <db>",
# "CREATE DATABASE <newdb> WITH TEMPLATE <db>"
return [{"type": "database"}]
elif token_v == "tableformat":
return [{"type": "table_format"}]
elif token_v.endswith(",") or is_operand(token_v) or token_v in ["=", "and", "or"]:
prev_keyword, text_before_cursor = find_prev_keyword(text_before_cursor)
if prev_keyword:
if token_v.endswith(","):
# handle edge case like mistake column name for keyword
# autocompleter.get_suggestions("select user, Passwor from mysql.user;", len("select user, Passwor"))
if prev_keyword and prev_keyword.value.lower() == "user":
parent_first_token = prev_keyword.parent.token_first()
if parent_first_token.value.lower() == "select":
return suggest_based_on_last_token(
"select", text_before_cursor, full_text, identifier
)
elif (
isinstance(parent_first_token.parent, IdentifierList)
and parent_first_token.parent.parent
):
parent_parent_first_token = (
parent_first_token.parent.parent.token_first()
)
if (
parent_parent_first_token
and parent_parent_first_token.value.lower() == "select"
):
return suggest_based_on_last_token(
"select", text_before_cursor, full_text, identifier
)
return suggest_based_on_last_token(
prev_keyword, text_before_cursor, full_text, identifier
)
else:
return []
else:
return [{"type": "keyword"}]
