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(): tables = extract_tables('select a,b from abc, def') assert sorted(tables) == [(None, 'abc', None), (None, 'def', 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_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_update_table(): tables = extract_tables('update abc set id = 1') assert tables == [(None, 'abc', None)]
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_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_simple_select_single_table(): tables = extract_tables('select * from abc') assert tables == [(None, 'abc', None)]
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"}]