diff options
Diffstat (limited to 'pgcli/packages')
| -rw-r--r-- | pgcli/packages/formatter/sqlformatter.py | 9 | ||||
| -rw-r--r-- | pgcli/packages/parseutils/__init__.py | 48 | ||||
| -rw-r--r-- | pgcli/packages/prompt_utils.py | 14 | ||||
| -rw-r--r-- | pgcli/packages/sqlcompletion.py | 3 |
4 files changed, 50 insertions, 24 deletions
diff --git a/pgcli/packages/formatter/sqlformatter.py b/pgcli/packages/formatter/sqlformatter.py index 5bf25fe..5224eff 100644 --- a/pgcli/packages/formatter/sqlformatter.py +++ b/pgcli/packages/formatter/sqlformatter.py @@ -14,10 +14,13 @@ preprocessors = () def escape_for_sql_statement(value): + if value is None: + return "NULL" + if isinstance(value, bytes): return f"X'{value.hex()}'" - else: - return "'{}'".format(value) + + return "'{}'".format(value) def adapter(data, headers, table_format=None, **kwargs): @@ -29,7 +32,7 @@ def adapter(data, headers, table_format=None, **kwargs): else: table_name = table[1] else: - table_name = '"DUAL"' + table_name = "DUAL" if table_format == "sql-insert": h = '", "'.join(headers) yield 'INSERT INTO "{}" ("{}") VALUES'.format(table_name, h) diff --git a/pgcli/packages/parseutils/__init__.py b/pgcli/packages/parseutils/__init__.py index 1acc008..023e13b 100644 --- a/pgcli/packages/parseutils/__init__.py +++ b/pgcli/packages/parseutils/__init__.py @@ -1,6 +1,17 @@ import sqlparse +BASE_KEYWORDS = [ + "drop", + "shutdown", + "delete", + "truncate", + "alter", + "unconditional_update", +] +ALL_KEYWORDS = BASE_KEYWORDS + ["update"] + + def query_starts_with(formatted_sql, prefixes): """Check if the query starts with any item from *prefixes*.""" prefixes = [prefix.lower() for prefix in prefixes] @@ -13,22 +24,35 @@ def query_is_unconditional_update(formatted_sql): return bool(tokens) and tokens[0] == "update" and "where" not in tokens -def query_is_simple_update(formatted_sql): - """Check if the query starts with UPDATE.""" - tokens = formatted_sql.split() - return bool(tokens) and tokens[0] == "update" - - -def is_destructive(queries, warning_level="all"): +def is_destructive(queries, keywords): """Returns if any of the queries in *queries* is destructive.""" - keywords = ("drop", "shutdown", "delete", "truncate", "alter") for query in sqlparse.split(queries): if query: formatted_sql = sqlparse.format(query.lower(), strip_comments=True).strip() - if query_starts_with(formatted_sql, keywords): - return True - if query_is_unconditional_update(formatted_sql): + if "unconditional_update" in keywords and query_is_unconditional_update( + formatted_sql + ): return True - if warning_level == "all" and query_is_simple_update(formatted_sql): + if query_starts_with(formatted_sql, keywords): return True return False + + +def parse_destructive_warning(warning_level): + """Converts a deprecated destructive warning option to a list of command keywords.""" + if not warning_level: + return [] + + if not isinstance(warning_level, list): + if "," in warning_level: + return warning_level.split(",") + warning_level = [warning_level] + + return { + "true": ALL_KEYWORDS, + "false": [], + "all": ALL_KEYWORDS, + "moderate": BASE_KEYWORDS, + "off": [], + "": [], + }.get(warning_level[0], warning_level) diff --git a/pgcli/packages/prompt_utils.py b/pgcli/packages/prompt_utils.py index e8589de..997b86e 100644 --- a/pgcli/packages/prompt_utils.py +++ b/pgcli/packages/prompt_utils.py @@ -3,7 +3,7 @@ import click from .parseutils import is_destructive -def confirm_destructive_query(queries, warning_level): +def confirm_destructive_query(queries, keywords, alias): """Check if the query is destructive and prompts the user to confirm. Returns: @@ -12,11 +12,13 @@ def confirm_destructive_query(queries, warning_level): * False if the query is destructive and the user doesn't want to proceed. """ - prompt_text = ( - "You're about to run a destructive command.\n" "Do you want to proceed? (y/n)" - ) - if is_destructive(queries, warning_level) and sys.stdin.isatty(): - return prompt(prompt_text, type=bool) + info = "You're about to run a destructive command" + if alias: + info += f" in {click.style(alias, fg='red')}" + + prompt_text = f"{info}.\nDo you want to proceed?" + if is_destructive(queries, keywords) and sys.stdin.isatty(): + return confirm(prompt_text) def confirm(*args, **kwargs): diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py index be4933a..b78edd6 100644 --- a/pgcli/packages/sqlcompletion.py +++ b/pgcli/packages/sqlcompletion.py @@ -290,7 +290,6 @@ def suggest_special(text): def suggest_based_on_last_token(token, stmt): - if isinstance(token, str): token_v = token.lower() elif isinstance(token, Comparison): @@ -399,7 +398,6 @@ def suggest_based_on_last_token(token, stmt): 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 @@ -436,7 +434,6 @@ def suggest_based_on_last_token(token, stmt): try: prev = stmt.get_previous_token(token).value.lower() if prev in ("drop", "alter", "create", "create or replace"): - # Suggest functions from either the currently-selected schema or the # public schema if no schema has been specified suggest = [] |