ported the desctructive warning from mycli

3 files changed, 47 insertions, 4 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index cfc3f2d7..b0096dc4 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -51,6 +51,7 @@ from .config import (get_casing_file,
 from .key_bindings import pgcli_bindings
 from .encodingutils import utf8tounicode
 from .encodingutils import text_type
+from .packages.prompt_utils import confirm_destructive_query
 from .__init__ import __version__
 
 click.disable_unicode_literals_warning = True
@@ -119,7 +120,7 @@ class PGCli(object):
     def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False,
                  pgexecute=None, pgclirc_file=None, row_limit=None,
                  single_connection=False, less_chatty=None, prompt=None, prompt_dsn=None,
-                 auto_vertical_output=False):
+                 auto_vertical_output=False, warn=None):
 
         self.force_passwd_prompt = force_passwd_prompt
         self.never_passwd_prompt = never_passwd_prompt
@@ -154,6 +155,8 @@ class PGCli(object):
         self.syntax_style = c['main']['syntax_style']
         self.cli_style = c['colors']
         self.wider_completion_menu = c['main'].as_bool('wider_completion_menu')
+        c_dest_warning = c['main'].as_bool('destructive_warning')
+        self.destructive_warning = c_dest_warning if warn is None else warn
         self.less_chatty = bool(less_chatty) or c['main'].as_bool('less_chatty')
         self.null_string = c['main'].get('null_string', '<null>')
         self.prompt_format = prompt if prompt is not None else c['main'].get('prompt', self.default_prompt)
@@ -280,6 +283,11 @@ class PGCli(object):
         except IOError as e:
             return [(None, None, None, str(e), '', False)]
 
+        if (self.destructive_warning and
+                confirm_destructive_query(query) is False):
+            message = 'Wise choice. Command execution stopped.'
+            return [(None, None, None, message)]
+
         on_error_resume = (self.on_error == 'RESUME')
         return self.pgexecute.run(
             query, self.pgspecial, on_error_resume=on_error_resume
@@ -460,7 +468,16 @@ class PGCli(object):
         logger = self.logger
 
         try:
-            output, query = self._evaluate_command(text)
+            if (self.destructive_warning):
+                destroy = confirm = confirm_destructive_query(text)
+                if destroy is None:
+                    output, query = self._evaluate_command(text)
+                elif destroy is True:
+                    click.secho('Your call!')
+                    output, query = self._evaluate_command(text)
+                else:
+                    click.secho('Wise choice!')
+                    raise KeyboardInterrupt
         except KeyboardInterrupt:
             # Restart connection to the database
             self.pgexecute.connect()
@@ -872,12 +889,14 @@ class PGCli(object):
               'available databases, then exit.')
 @click.option('--auto-vertical-output', is_flag=True,
               help='Automatically switch to vertical output mode if the result is wider than the terminal width.')
+@click.option('--warn/--no-warn', default=None,
+              help='Warn before running a destructive query.')
 @click.argument('database', default=lambda: None, envvar='PGDATABASE', nargs=1)
 @click.argument('username', default=lambda: None, envvar='PGUSER', nargs=1)
 def cli(database, username_opt, host, port, prompt_passwd, never_prompt,
         single_connection, dbname, username, version, pgclirc, dsn, row_limit,
         less_chatty, prompt, prompt_dsn, list_databases, auto_vertical_output,
-        list_dsn):
+        list_dsn, warn):
 
     if version:
         print('Version:', __version__)
@@ -913,7 +932,7 @@ def cli(database, username_opt, host, port, prompt_passwd, never_prompt,
     pgcli = PGCli(prompt_passwd, never_prompt, pgclirc_file=pgclirc,
                   row_limit=row_limit, single_connection=single_connection,
                   less_chatty=less_chatty, prompt=prompt, prompt_dsn=prompt_dsn,
-                  auto_vertical_output=auto_vertical_output)
+                  auto_vertical_output=auto_vertical_output, warn=warn)
 
     # Choose which ever one has a valid value.
     database = database or dbname
diff --git a/pgcli/packages/parseutils/__init__.py b/pgcli/packages/parseutils/__init__.py
index e69de29b..818af9a7 100644
--- a/pgcli/packages/parseutils/__init__.py
+++ b/pgcli/packages/parseutils/__init__.py
@@ -0,0 +1,19 @@
+import sqlparse
+
+def query_starts_with(query, prefixes):
+    """Check if the query starts with any item from *prefixes*."""
+    prefixes = [prefix.lower() for prefix in prefixes]
+    formatted_sql = sqlparse.format(query.lower(), strip_comments=True)
+    return bool(formatted_sql) and formatted_sql.split()[0] in prefixes
+
+def queries_start_with(queries, prefixes):
+    """Check if any queries start with any item from *prefixes*."""
+    for query in sqlparse.split(queries):
+        if query and query_starts_with(query, prefixes) is True:
+            return True
+    return False
+
+def is_destructive(queries):
+    """Returns if any of the queries in *queries* is destructive."""
+    keywords = ('drop', 'shutdown', 'delete', 'truncate', 'alter')
+    return queries_start_with(queries, keywords)
diff --git a/pgcli/pgclirc b/pgcli/pgclirc
index 70f10ee7..f2061c71 100644
--- a/pgcli/pgclirc
+++ b/pgcli/pgclirc
@@ -22,6 +22,11 @@ multi_line = False
 # a command.
 multi_line_mode = psql
 
+# Destructive warning mode will alert you before executing a sql statement
+# that may cause harm to the database such as "drop table", "drop database"
+# or "shutdown".
+destructive_warning = True
+
 # Enables expand mode, which is similar to `\x` in psql.
 expand = False