@@ -53,7 +53,7 @@ def __init__(self, url, **kwargs):
53
53
import sqlalchemy
54
54
import sqlalchemy .orm
55
55
import threading
56
-
56
+
57
57
# Temporary fix for missing sqlite3 module on the buildpack stack
58
58
try :
59
59
import sqlite3
@@ -149,15 +149,15 @@ def execute(self, sql, *args, **kwargs):
149
149
if len (args ) > 0 and len (kwargs ) > 0 :
150
150
raise RuntimeError ("cannot pass both positional and named parameters" )
151
151
152
- # Infer command from (unflattened) statement
153
- for token in statements [0 ]:
154
- if token . ttype in [ sqlparse . tokens . Keyword , sqlparse . tokens . Keyword . DDL , sqlparse . tokens . Keyword . DML ]:
155
- token_value = token . value . upper ()
156
- if token_value in [ "BEGIN" , "DELETE" , "INSERT" , "SELECT" , "START" , "UPDATE" ]:
157
- command = token_value
158
- break
159
- else :
160
- command = None
152
+ # Infer command from flattened statement to a single string separated by spaces
153
+ full_statement = ' ' . join ( str ( token ) for token in statements [0 ]. tokens if token . ttype in [ sqlparse . tokens . Keyword , sqlparse . tokens . Keyword . DDL , sqlparse . tokens . Keyword . DML ])
154
+ full_statement = full_statement . upper ()
155
+
156
+ # set of possible commands
157
+ commands = { "BEGIN" , "CREATE VIEW" , "DELETE" , "INSERT" , "SELECT" , "START" , "UPDATE" }
158
+
159
+ # check if the full_statement starts with any command
160
+ command = next (( cmd for cmd in commands if full_statement . startswith ( cmd )), None )
161
161
162
162
# Flatten statement
163
163
tokens = list (statements [0 ].flatten ())
@@ -393,6 +393,10 @@ def teardown_appcontext(exception):
393
393
elif command in ["DELETE" , "UPDATE" ]:
394
394
ret = result .rowcount
395
395
396
+ # If CREATE VIEW, return True
397
+ elif command == "CREATE VIEW" :
398
+ ret = True
399
+
396
400
# If constraint violated
397
401
except sqlalchemy .exc .IntegrityError as e :
398
402
self ._logger .error (termcolor .colored (_statement , "red" ))
0 commit comments