Finer control over destructive warning. (#1242)

* Finer control over destructive warning.

* Review feedback.

* Changelog.

* Run integration tests with --warn=moderate.

* Fix typo.

* Black.

10 files changed, 103 insertions, 50 deletions

diff --git a/README.rst b/README.rst
index d5934274..1fa4e7f5 100644
--- a/README.rst
+++ b/README.rst
@@ -62,32 +62,32 @@ For more details:
     Usage: pgcli [OPTIONS] [DBNAME] [USERNAME]
 
     Options:
-      -h, --host TEXT         Host address of the postgres database.
-      -p, --port INTEGER      Port number at which the postgres instance is
-                              listening.
-      -U, --username TEXT     Username to connect to the postgres database.
-      -u, --user TEXT         Username to connect to the postgres database.
-      -W, --password          Force password prompt.
-      -w, --no-password       Never prompt for password.
-      --single-connection     Do not use a separate connection for completions.
-      -v, --version           Version of pgcli.
-      -d, --dbname TEXT       database name to connect to.
-      --pgclirc PATH          Location of pgclirc file.
-      -D, --dsn TEXT          Use DSN configured into the [alias_dsn] section of
-                              pgclirc file.
-      --list-dsn              list of DSN configured into the [alias_dsn] section
-                              of pgclirc file.
-      --row-limit INTEGER     Set threshold for row limit prompt. Use 0 to disable
-                              prompt.
-      --less-chatty           Skip intro on startup and goodbye on exit.
-      --prompt TEXT           Prompt format (Default: "\u@\h:\d> ").
-      --prompt-dsn TEXT       Prompt format for connections using DSN aliases
-                              (Default: "\u@\h:\d> ").
-      -l, --list              list available databases, then exit.
-      --auto-vertical-output  Automatically switch to vertical output mode if the
-                              result is wider than the terminal width.
-      --warn / --no-warn      Warn before running a destructive query.
-      --help                  Show this message and exit.
+      -h, --host TEXT            Host address of the postgres database.
+      -p, --port INTEGER         Port number at which the postgres instance is
+                                 listening.
+      -U, --username TEXT        Username to connect to the postgres database.
+      -u, --user TEXT            Username to connect to the postgres database.
+      -W, --password             Force password prompt.
+      -w, --no-password          Never prompt for password.
+      --single-connection        Do not use a separate connection for completions.
+      -v, --version              Version of pgcli.
+      -d, --dbname TEXT          database name to connect to.
+      --pgclirc FILE             Location of pgclirc file.
+      -D, --dsn TEXT             Use DSN configured into the [alias_dsn] section
+                                 of pgclirc file.
+      --list-dsn                 list of DSN configured into the [alias_dsn]
+                                 section of pgclirc file.
+      --row-limit INTEGER        Set threshold for row limit prompt. Use 0 to
+                                 disable prompt.
+      --less-chatty              Skip intro on startup and goodbye on exit.
+      --prompt TEXT              Prompt format (Default: "\u@\h:\d> ").
+      --prompt-dsn TEXT          Prompt format for connections using DSN aliases
+                                 (Default: "\u@\h:\d> ").
+      -l, --list                 list available databases, then exit.
+      --auto-vertical-output     Automatically switch to vertical output mode if
+                                 the result is wider than the terminal width.
+      --warn [all|moderate|off]  Warn before running a destructive query.
+      --help                     Show this message and exit.
 
 ``pgcli`` also supports many of the same `environment variables`_ as ``psql`` for login options (e.g. ``PGHOST``, ``PGPORT``, ``PGUSER``, ``PGPASSWORD``, ``PGDATABASE``).
 
diff --git a/changelog.rst b/changelog.rst
index 224934c1..6ae26ca5 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -1,6 +1,12 @@
 TBD
 ===
 
+Features:
+---------
+
+* Consider `update` queries destructive and issue a warning. Change
+  `destructive_warning` setting to `all|moderate|off`, vs `true|false`. (#1239)
+
 Bug fixes:
 ----------
 
diff --git a/pgcli/main.py b/pgcli/main.py
index 7ab973a2..539ba449 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -201,8 +201,11 @@ class PGCli:
         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.destructive_warning = warn or c["main"]["destructive_warning"]
+        # also handle boolean format of destructive warning
+        self.destructive_warning = {"true": "all", "false": "off"}.get(
+            self.destructive_warning.lower(), self.destructive_warning
+        )
         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 = (
@@ -389,7 +392,10 @@ class PGCli:
         except OSError as e:
             return [(None, None, None, str(e), "", False, True)]
 
-        if self.destructive_warning and confirm_destructive_query(query) is False:
+        if (
+            self.destructive_warning != "off"
+            and confirm_destructive_query(query, self.destructive_warning) is False
+        ):
             message = "Wise choice. Command execution stopped."
             return [(None, None, None, message)]
 
@@ -644,8 +650,10 @@ class PGCli:
         query = MetaQuery(query=text, successful=False)
 
         try:
-            if self.destructive_warning:
-                destroy = confirm = confirm_destructive_query(text)
+            if self.destructive_warning != "off":
+                destroy = confirm = confirm_destructive_query(
+                    text, self.destructive_warning
+                )
                 if destroy is False:
                     click.secho("Wise choice!")
                     raise KeyboardInterrupt
@@ -1188,7 +1196,10 @@ class PGCli:
     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."
+    "--warn",
+    default=None,
+    type=click.Choice(["all", "moderate", "off"]),
+    help="Warn before running a destructive query.",
 )
 @click.argument("dbname", default=lambda: None, envvar="PGDATABASE", nargs=1)
 @click.argument("username", default=lambda: None, envvar="PGUSER", nargs=1)
diff --git a/pgcli/packages/parseutils/__init__.py b/pgcli/packages/parseutils/__init__.py
index a11e7bf2..1acc008e 100644
--- a/pgcli/packages/parseutils/__init__.py
+++ b/pgcli/packages/parseutils/__init__.py
@@ -1,22 +1,34 @@
 import sqlparse
 
 
-def query_starts_with(query, prefixes):
+def query_starts_with(formatted_sql, 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).strip()
     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 query_is_unconditional_update(formatted_sql):
+    """Check if the query starts with UPDATE and contains no WHERE."""
+    tokens = formatted_sql.split()
+    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):
+
+def is_destructive(queries, warning_level="all"):
     """Returns if any of the queries in *queries* is destructive."""
     keywords = ("drop", "shutdown", "delete", "truncate", "alter")
-    return queries_start_with(queries, keywords)
+    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):
+                return True
+            if warning_level == "all" and query_is_simple_update(formatted_sql):
+                return True
+    return False
diff --git a/pgcli/packages/prompt_utils.py b/pgcli/packages/prompt_utils.py
index 3c584908..e8589def 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):
+def confirm_destructive_query(queries, warning_level):
     """Check if the query is destructive and prompts the user to confirm.
 
     Returns:
@@ -15,7 +15,7 @@ def confirm_destructive_query(queries):
     prompt_text = (
         "You're about to run a destructive command.\n" "Do you want to proceed? (y/n)"
     )
-    if is_destructive(queries) and sys.stdin.isatty():
+    if is_destructive(queries, warning_level) and sys.stdin.isatty():
         return prompt(prompt_text, type=bool)
 
 
diff --git a/pgcli/pgclirc b/pgcli/pgclirc
index d9a8c2b4..15c10f5e 100644
--- a/pgcli/pgclirc
+++ b/pgcli/pgclirc
@@ -23,9 +23,13 @@ multi_line = False
 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
+# that may cause harm to the database such as "drop table", "drop database",
+# "shutdown", "delete", or "update".
+# Possible values:
+# "all" - warn on data definition statements, server actions such as SHUTDOWN, DELETE or UPDATE
+# "moderate" - skip warning on UPDATE statements, except for unconditional updates
+# "off" - skip all warnings
+destructive_warning = all
 
 # Enables expand mode, which is similar to `\x` in psql.
 expand = False
diff --git a/tests/conftest.py b/tests/conftest.py
index 2a715b11..33cddf24 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -12,7 +12,7 @@ from utils import (
 import pgcli.pgexecute
 
 
-@pytest.yield_fixture(scope="function")
+@pytest.fixture(scope="function")
 def connection():
     create_db("_test_db")
     connection = db_connection("_test_db")
diff --git a/tests/features/environment.py b/tests/features/environment.py
index a8b3e26a..215c85cd 100644
--- a/tests/features/environment.py
+++ b/tests/features/environment.py
@@ -63,7 +63,7 @@ def before_all(context):
                         "import coverage",
                         "coverage.process_startup()",
                         "import pgcli.main",
-                        "pgcli.main.cli()",
+                        "pgcli.main.cli(auto_envvar_prefix='BEHAVE')",
                     ]
                 ),
             )
@@ -102,6 +102,7 @@ def before_all(context):
     else:
         if "PGPASSWORD" in os.environ:
             del os.environ["PGPASSWORD"]
+    os.environ["BEHAVE_WARN"] = "moderate"
 
     context.cn = dbutils.create_db(
         context.conf["host"],
diff --git a/tests/parseutils/test_parseutils.py b/tests/parseutils/test_parseutils.py
index aaeb3f91..5a375d70 100644
--- a/tests/parseutils/test_parseutils.py
+++ b/tests/parseutils/test_parseutils.py
@@ -1,4 +1,5 @@
 import pytest
+from pgcli.packages.parseutils import is_destructive
 from pgcli.packages.parseutils.tables import extract_tables
 from pgcli.packages.parseutils.utils import find_prev_keyword, is_open_quote
 
@@ -259,3 +260,21 @@ def test_is_open_quote__closed(sql):
 )
 def test_is_open_quote__open(sql):
     assert is_open_quote(sql)
+
+
+@pytest.mark.parametrize(
+    ("sql", "warning_level", "expected"),
+    [
+        ("update abc set x = 1", "all", True),
+        ("update abc set x = 1 where y = 2", "all", True),
+        ("update abc set x = 1", "moderate", True),
+        ("update abc set x = 1 where y = 2", "moderate", False),
+        ("select x, y, z from abc", "all", False),
+        ("drop abc", "all", True),
+        ("alter abc", "all", True),
+        ("delete abc", "all", True),
+        ("truncate abc", "all", True),
+    ],
+)
+def test_is_destructive(sql, warning_level, expected):
+    assert is_destructive(sql, warning_level=warning_level) == expected
diff --git a/tests/test_prompt_utils.py b/tests/test_prompt_utils.py
index c1f8a169..a8a3a1e0 100644
--- a/tests/test_prompt_utils.py
+++ b/tests/test_prompt_utils.py
@@ -7,4 +7,4 @@ def test_confirm_destructive_query_notty():
     stdin = click.get_text_stream("stdin")
     if not stdin.isatty():
         sql = "drop database foo;"
-        assert confirm_destructive_query(sql) is None
+        assert confirm_destructive_query(sql, "all") is None