diff options
Diffstat (limited to 'pgcli/packages')
-rw-r--r-- | pgcli/packages/formatter/__init__.py | 1 | ||||
-rw-r--r-- | pgcli/packages/formatter/sqlformatter.py | 71 | ||||
-rw-r--r-- | pgcli/packages/parseutils/tables.py | 2 | ||||
-rw-r--r-- | pgcli/packages/sqlcompletion.py | 2 |
4 files changed, 74 insertions, 2 deletions
diff --git a/pgcli/packages/formatter/__init__.py b/pgcli/packages/formatter/__init__.py new file mode 100644 index 0000000..9bad579 --- /dev/null +++ b/pgcli/packages/formatter/__init__.py @@ -0,0 +1 @@ +# coding=utf-8 diff --git a/pgcli/packages/formatter/sqlformatter.py b/pgcli/packages/formatter/sqlformatter.py new file mode 100644 index 0000000..5bf25fe --- /dev/null +++ b/pgcli/packages/formatter/sqlformatter.py @@ -0,0 +1,71 @@ +# coding=utf-8 + +from pgcli.packages.parseutils.tables import extract_tables + + +supported_formats = ( + "sql-insert", + "sql-update", + "sql-update-1", + "sql-update-2", +) + +preprocessors = () + + +def escape_for_sql_statement(value): + if isinstance(value, bytes): + return f"X'{value.hex()}'" + else: + return "'{}'".format(value) + + +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 register_new_formatter(TabularOutputFormatter): + global formatter + formatter = TabularOutputFormatter + for sql_format in supported_formats: + TabularOutputFormatter.register_new_formatter( + sql_format, adapter, preprocessors, {"table_format": sql_format} + ) diff --git a/pgcli/packages/parseutils/tables.py b/pgcli/packages/parseutils/tables.py index f2e1e42..9098115 100644 --- a/pgcli/packages/parseutils/tables.py +++ b/pgcli/packages/parseutils/tables.py @@ -139,7 +139,7 @@ def extract_table_identifiers(token_stream, allow_functions=True): # extract_tables is inspired from examples in the sqlparse lib. def extract_tables(sql): - """Extract the table names from an SQL statment. + """Extract the table names from an SQL statement. Returns a list of TableReference namedtuples diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py index 6305301..be4933a 100644 --- a/pgcli/packages/sqlcompletion.py +++ b/pgcli/packages/sqlcompletion.py @@ -380,7 +380,7 @@ def suggest_based_on_last_token(token, stmt): ) elif p.token_first().value.lower() == "select": - # If the lparen is preceeded by a space chances are we're about to + # If the lparen is preceded 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(),) |