Merge pull request #796 from SukiCZ/fix/use_after_create_delete

Fix/use after create delete

Author: tsroten

changelog.md

@@ -16,6 +16,7 @@ Features:
 * Mark `update` without `where`-clause as destructive query (Thanks: [Klaus Wünschel]).
 * Added DELIMITER command (Thanks: [Georgy Frolov])
 * Added clearer error message when failing to connect to the default socket.
+* Extend main.is_dropping_database check with create after delete statement.
 * Search `${XDG_CONFIG_HOME}/mycli/myclirc` after `${HOME}/.myclirc` and before `/etc/myclirc` (Thanks: [Takeshi D. Itoh])
 
 Bug Fixes:

mycli/AUTHORS

@@ -68,6 +68,7 @@ Contributors:
   * Mike Palandra
   * Georgy Frolov
   * Jonathan Lloyd
+  * Jakub Boukal
   * Takeshi D. Itoh
   * laixintao
 

mycli/packages/parseutils.py

@@ -241,6 +241,7 @@ def is_open_quote(sql):
 
 def is_dropping_database(queries, dbname):
     """Determine if the query is dropping a specific database."""
+    result = False
     if dbname is None:
         return False
 
@@ -253,16 +254,14 @@ def normalize_db_name(db):
         keywords = [t for t in query.tokens if t.is_keyword]
         if len(keywords) < 2:
             continue
-        if keywords[0].normalized == "DROP" and keywords[1].value.lower() in (
+        if keywords[0].normalized in ("DROP", "CREATE") and keywords[1].value.lower() in (
             "database",
             "schema",
         ):
             database_token = next(
                 (t for t in query.tokens if isinstance(t, Identifier)), None
             )
-            return (
-                database_token is not None
-                and normalize_db_name(database_token.get_name()) == dbname
-            )
+            if database_token is not None and normalize_db_name(database_token.get_name()) == dbname:
+                result = keywords[0].normalized == "DROP"
     else:
-        return False
+        return result

test/test_parseutils.py

@@ -175,6 +175,8 @@ def test_query_has_where_clause(sql, has_where_clause):
         ('drop schema foo', 'bar', False),
         ('drop database bar', 'foo', False),
         ('drop database foo', None, False),
+        ('drop database foo; create database foo', 'foo', False),
+        ('drop database foo; create database bar', 'foo', True),
         ('select bar from foo; drop database bazz', 'foo', False),
         ('select bar from foo; drop database bazz', 'bazz', True),
         ('-- dropping database \n '