fixed support for transactions

Author: dmalan

src/cs50/sql.py

@@ -56,13 +56,14 @@ def __init__(self, url, **kwargs):
             if not os.path.isfile(matches.group(1)):
                 raise RuntimeError("not a file: {}".format(matches.group(1)))
 
-        # Create engine, raising exception if back end's module not installed
-        self._engine = sqlalchemy.create_engine(url, **kwargs).execution_options(autocommit=True)
+        # Create engine, disabling SQLAlchemy's own autocommit mode, raising exception if back end's module not installed
+        self._engine = sqlalchemy.create_engine(url, **kwargs).execution_options(autocommit=False)
 
         # Listener for connections
         def connect(dbapi_connection, connection_record):
 
-            # Disable underlying API's own emitting of BEGIN and COMMIT
+            # Disable underlying API's own emitting of BEGIN and COMMIT so we can ourselves
+            # https://docs.sqlalchemy.org/en/13/dialects/sqlite.html#serializable-isolation-savepoints-transactional-ddl
             dbapi_connection.isolation_level = None
 
             # Enable foreign key constraints
@@ -71,6 +72,9 @@ def connect(dbapi_connection, connection_record):
                 cursor.execute("PRAGMA foreign_keys=ON")
                 cursor.close()
 
+            # Autocommit by default
+            self._autocommit = True
+
         # Register listener
         sqlalchemy.event.listen(self._engine, "connect", connect)
 
@@ -90,9 +94,14 @@ def connect(dbapi_connection, connection_record):
             self._logger.disabled = disabled
 
     def __del__(self):
+        """Disconnect from database."""
+        self._disconnect()
+
+    def _disconnect(self):
         """Close database connection."""
         if hasattr(self, "_connection"):
             self._connection.close()
+            delattr(self, "_connection")
 
     @_enable_logging
     def execute(self, sql, *args, **kwargs):
@@ -107,7 +116,7 @@ def execute(self, sql, *args, **kwargs):
         import warnings
 
         # Parse statement, stripping comments and then leading/trailing whitespace
-        statements = sqlparse.parse(sqlparse.format(sql, strip_comments=True).strip())
+        statements = sqlparse.parse(sqlparse.format(sql, keyword_case="upper", strip_comments=True).strip())
 
         # Allow only one statement at a time, since SQLite doesn't support multiple
         # https://docs.python.org/3/library/sqlite3.html#sqlite3.Cursor.execute
@@ -122,9 +131,10 @@ def execute(self, sql, *args, **kwargs):
 
         # Infer command from (unflattened) statement
         for token in statements[0]:
-            if token.ttype in [sqlparse.tokens.Keyword.DDL, sqlparse.tokens.Keyword.DML]:
-                command = token.value.upper()
-                break
+            if token.ttype in [sqlparse.tokens.Keyword, sqlparse.tokens.Keyword.DDL, sqlparse.tokens.Keyword.DML]:
+                if token.value in ["BEGIN", "DELETE", "INSERT", "SELECT", "START", "UPDATE"]:
+                    command = token.value
+                    break
         else:
             command = None
 
@@ -316,8 +326,21 @@ def shutdown_session(exception=None):
                 # Join tokens into statement, abbreviating binary data as <class 'bytes'>
                 _statement = "".join([str(bytes) if token.ttype == sqlparse.tokens.Other else str(token) for token in tokens])
 
+                # Check for start of transaction
+                if command in ["BEGIN", "START"]:
+                    self._autocommit = False
+
                 # Execute statement
-                result = connection.execute(sqlalchemy.text(statement))
+                if self._autocommit:
+                    connection.execute(sqlalchemy.text("BEGIN"))
+                    result = connection.execute(sqlalchemy.text(statement))
+                    connection.execute(sqlalchemy.text("COMMIT"))
+                else:
+                    result = connection.execute(sqlalchemy.text(statement))
+
+                # Check for end of transaction
+                if command in ["COMMIT", "ROLLBACK"]:
+                    self._autocommit = True
 
                 # Return value
                 ret = True
@@ -359,13 +382,15 @@ def shutdown_session(exception=None):
 
             # If constraint violated, return None
             except sqlalchemy.exc.IntegrityError as e:
+                self._disconnect()
                 self._logger.debug(termcolor.colored(statement, "yellow"))
                 e = RuntimeError(e.orig)
                 e.__cause__ = None
                 raise e
 
             # If user errror
             except sqlalchemy.exc.OperationalError as e:
+                self._disconnect()
                 self._logger.debug(termcolor.colored(statement, "red"))
                 e = RuntimeError(e.orig)
                 e.__cause__ = None

tests/foo.py

@@ -5,11 +5,11 @@
 
 import cs50
 
+"""
 db = cs50.SQL("sqlite:///foo.db")
 
 logging.getLogger("cs50").disabled = False
 
-"""
 #db.execute("SELECT ? FROM ? ORDER BY ?", "a", "tbl", "c")
 db.execute("CREATE TABLE IF NOT EXISTS bar (firstname STRING)")
 
@@ -36,16 +36,13 @@
 print(db.execute("DROP TABLE IF EXISTS cs50"))
 print(db.execute("CREATE TABLE cs50 (id SERIAL PRIMARY KEY, val VARCHAR(16), bin BYTEA)"))
 print(db.execute("INSERT INTO cs50 (val) VALUES('foo')"))
-print(db.execute("INSERT INTO cs50 (val) VALUES('foo')"))
-print(db.execute("INSERT INTO cs50 (val) VALUES('foo')"))
+print(db.execute("INSERT INTO cs50 (val) VALUES('bar')"))
+print(db.execute("INSERT INTO cs50 (val) VALUES('baz')"))
 print(db.execute("SELECT * FROM cs50"))
-print(db.execute("COMMIT"))
-"""
 try:
     print(db.execute("INSERT INTO cs50 (id, val) VALUES(1, 'bar')"))
 except Exception as e:
     print(e)
     pass
-print(db.execute("INSERT INTO cs50 (val) VALUES('foo')"))
-print(db.execute("DELETE FROM cs50"))
-"""
+print(db.execute("INSERT INTO cs50 (val) VALUES('qux')"))
+#print(db.execute("DELETE FROM cs50"))