Use sessions to handle transactions, allowing for both auto and manual

commit modes. Registers Flask appcontext teardown function only once per
database instance, and also allows for multiple database connections in
a single Flask request.

Add unit tests for SQL savepoints, autocommit mode, manual transaction
mode. Add integration tests for Flask.

Author: jsarchibald

src/cs50/sql.py

@@ -43,6 +43,7 @@ def __init__(self, url, **kwargs):
         import os
         import re
         import sqlalchemy
+        import sqlalchemy.orm as orm
         import sqlite3
 
         # Get logger
@@ -56,9 +57,16 @@ def __init__(self, url, **kwargs):
             if not os.path.isfile(matches.group(1)):
                 raise RuntimeError("not a file: {}".format(matches.group(1)))
 
+        # Record the URL (used in testing)
+        self.url = url
+
         # 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)
 
+        # Create a variable to hold the session. If None, autocommit is on.
+        self.Session = orm.sessionmaker(bind=self._engine)
+        self._session = None
+
         # Listener for connections
         def connect(dbapi_connection, connection_record):
 
@@ -90,9 +98,9 @@ def connect(dbapi_connection, connection_record):
             self._logger.disabled = disabled
 
     def __del__(self):
-        """Close database connection."""
-        if hasattr(self, "_connection"):
-            self._connection.close()
+        """Close database session and connection."""
+        if self._session is not None:
+            self._session.close()
 
     @_enable_logging
     def execute(self, sql, *args, **kwargs):
@@ -125,6 +133,12 @@ def execute(self, sql, *args, **kwargs):
             if token.ttype in [sqlparse.tokens.Keyword.DDL, sqlparse.tokens.Keyword.DML]:
                 command = token.value.upper()
                 break
+
+            # Begin a new transaction session, if done manually
+            elif token.value.upper() in ["BEGIN", "START"]:
+                if self._session is not None:
+                    self._session.close()
+                self._session = self.Session()
         else:
             command = None
 
@@ -272,6 +286,11 @@ def execute(self, sql, *args, **kwargs):
         statement = "".join([str(token) for token in tokens])
 
         # Connect to database (for transactions' sake)
+        session = self._session
+        if session is None:
+            session = self.Session()
+
+        # Set up a Flask app teardown function to close session at teardown
         try:
 
             # Infer whether Flask is installed
@@ -280,29 +299,18 @@ def execute(self, sql, *args, **kwargs):
             # Infer whether app is defined
             assert flask.current_app
 
-            # If no connection for app's current request yet
-            if not hasattr(flask.g, "_connection"):
-
-                # Connect now
-                flask.g._connection = self._engine.connect()
+            # Disconnect later - but only once
+            if not hasattr(self, "teardown_appcontext_added"):
+                self.teardown_appcontext_added = True
 
-                # Disconnect later
+                # Register shutdown_session on app context teardown
                 @flask.current_app.teardown_appcontext
                 def shutdown_session(exception=None):
-                    if hasattr(flask.g, "_connection"):
-                        flask.g._connection.close()
-
-            # Use this connection
-            connection = flask.g._connection
+                    if self._session is not None:
+                        self._session.close()
 
         except (ModuleNotFoundError, AssertionError):
-
-            # If no connection yet
-            if not hasattr(self, "_connection"):
-                self._connection = self._engine.connect()
-
-            # Use this connection
-            connection = self._connection
+            pass
 
         # Catch SQLAlchemy warnings
         with warnings.catch_warnings():
@@ -317,7 +325,7 @@ def shutdown_session(exception=None):
                 _statement = "".join([str(bytes) if token.ttype == sqlparse.tokens.Other else str(token) for token in tokens])
 
                 # Execute statement
-                result = connection.execute(sqlalchemy.text(statement))
+                result = session.execute(sqlalchemy.text(statement))
 
                 # Return value
                 ret = True
@@ -346,7 +354,7 @@ def shutdown_session(exception=None):
                 elif command == "INSERT":
                     if self._engine.url.get_backend_name() in ["postgres", "postgresql"]:
                         try:
-                            result = connection.execute("SELECT LASTVAL()")
+                            result = session.execute("SELECT LASTVAL()")
                             ret = result.first()[0]
                         except sqlalchemy.exc.OperationalError:  # If lastval is not yet defined in this session
                             ret = None
@@ -357,6 +365,18 @@ def shutdown_session(exception=None):
                 elif command in ["DELETE", "UPDATE"]:
                     ret = result.rowcount
 
+                # If COMMIT or ROLLBACK, turn on autocommit mode
+                elif command in ["COMMIT", "ROLLBACK"] and "TO" not in statement:
+                    session.close()
+                    self._session = None
+
+
+                # If autocommit is on, commit and close
+                if self._session is None and command not in ["COMMIT", "ROLLBACK"]:
+                    if command not in ["SELECT"]:
+                        session.commit()
+                    session.close()
+
             # If constraint violated, return None
             except sqlalchemy.exc.IntegrityError as e:
                 self._logger.debug(termcolor.colored(statement, "yellow"))

tests/flask/application.py

@@ -1,3 +1,5 @@
+import logging
+import os
 import requests
 import sys
 from flask import Flask, render_template
@@ -9,14 +11,64 @@
 
 app = Flask(__name__)
 
-db = cs50.SQL("sqlite:///../sqlite.db")
+logging.disable(logging.CRITICAL)
+os.environ["WERKZEUG_RUN_MAIN"] = "true"
+
+db = cs50.SQL("sqlite:///../test.db")
 
 @app.route("/")
 def index():
-    db.execute("SELECT 1")
     """
     def f():
         res = requests.get("cs50.harvard.edu")
     f()
     """
     return render_template("index.html")
+
+@app.route("/autocommit")
+def autocommit():
+    db.execute("INSERT INTO test (val) VALUES (?)", "def")
+    db2 = cs50.SQL(db.url)
+    ret = db2.execute("SELECT val FROM test WHERE val=?", "def")
+    return str(ret == [{"val": "def"}])
+
+@app.route("/create")
+def create():
+    ret = db.execute("CREATE TABLE test (id INTEGER PRIMARY KEY AUTOINCREMENT, val VARCHAR(16))")
+    return str(ret)
+
+@app.route("/delete")
+def delete():
+    ret = db.execute("DELETE FROM test")
+    return str(ret > 0)
+
+@app.route("/drop")
+def drop():
+    ret = db.execute("DROP TABLE test")
+    return str(ret)
+
+@app.route("/insert")
+def insert():
+    ret = db.execute("INSERT INTO test (val) VALUES (?)", "abc")
+    return str(ret > 0)
+
+@app.route("/multiple_connections")
+def multiple_connections():
+    ctx = len(app.teardown_appcontext_funcs)
+    db1 = cs50.SQL(db.url)
+    td1 = (len(app.teardown_appcontext_funcs) == ctx + 1)
+    db2 = cs50.SQL(db.url)
+    td2 = (len(app.teardown_appcontext_funcs) == ctx + 2)
+    return str(td1 and td2)
+
+@app.route("/select")
+def select():
+    ret = db.execute("SELECT val FROM test")
+    return str(ret == [{"val": "abc"}])
+
+@app.route("/single_teardown")
+def single_teardown():
+    db.execute("SELECT * FROM test")
+    ctx = len(app.teardown_appcontext_funcs)
+    db.execute("SELECT COUNT(id) FROM test")
+    return str(ctx == len(app.teardown_appcontext_funcs))

tests/flask/test.py

@@ -0,0 +1,49 @@
+from application import app
+import logging
+import requests
+import sys
+import threading
+import time
+import unittest
+
+
+def request(route):
+    r = requests.get("http://localhost:5000/{}".format(route))
+    return r.text == "True"
+
+class FlaskTests(unittest.TestCase):
+
+    def test__create(self):
+        self.assertTrue(request("create"))
+ 
+    def test_autocommit(self):
+        self.assertTrue(request("autocommit"))
+
+    def test_delete(self):
+        self.assertTrue(request("delete"))
+
+    def test_insert(self):
+        self.assertTrue(request("insert"))
+
+    def test_multiple_connections(self):
+        self.assertTrue(request("multiple_connections"))
+
+    def test_select(self):
+        self.assertTrue(request("select"))
+
+    def test_single_teardown(self):
+        self.assertTrue(request("single_teardown"))
+
+    def test_zdrop(self):
+        self.assertTrue(request("drop"))
+
+
+if __name__ == "__main__":
+    t = threading.Thread(target=app.run, daemon=True)
+    t.start()
+
+    suite = unittest.TestSuite([
+        unittest.TestLoader().loadTestsFromTestCase(FlaskTests)
+    ])
+
+    sys.exit(not unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful())

tests/sql.py

@@ -115,11 +115,22 @@ def test_blob(self):
             self.db.execute("INSERT INTO cs50(bin) VALUES(:bin)", bin=row["bin"])
         self.assertEqual(self.db.execute("SELECT id, bin FROM cs50"), rows)
 
+    def test_autocommit(self):
+        self.assertEqual(self.db.execute("INSERT INTO cs50(val) VALUES('foo')"), 1)
+        self.assertEqual(self.db.execute("INSERT INTO cs50(val) VALUES('bar')"), 2)
+
+        # Load a new database instance to confirm the INSERTs were committed
+        db2 = SQL(self.db.url)
+        self.assertEqual(db2.execute("DELETE FROM cs50 WHERE id < 3"), 2)
+
     def test_commit(self):
         self.db.execute("BEGIN")
         self.db.execute("INSERT INTO cs50 (val) VALUES('foo')")
         self.db.execute("COMMIT")
-        self.assertEqual(self.db.execute("SELECT val FROM cs50"), [{"val": "foo"}])
+
+        # Load a new database instance to confirm the INSERT was committed
+        db2 = SQL(self.db.url)
+        self.assertEqual(db2.execute("SELECT val FROM cs50"), [{"val": "foo"}])
 
     def test_rollback(self):
         self.db.execute("BEGIN")
@@ -128,6 +139,17 @@ def test_rollback(self):
         self.db.execute("ROLLBACK")
         self.assertEqual(self.db.execute("SELECT val FROM cs50"), [])
 
+    def test_savepoint(self):
+        self.db.execute("BEGIN")
+        self.db.execute("INSERT INTO cs50 (val) VALUES('foo')")
+        self.db.execute("SAVEPOINT sp1")
+        self.db.execute("INSERT INTO cs50 (val) VALUES('bar')")
+        self.assertEqual(self.db.execute("SELECT val FROM cs50"), [{"val": "foo"}, {"val": "bar"}])
+        self.db.execute("ROLLBACK TO sp1")
+        self.assertEqual(self.db.execute("SELECT val FROM cs50"), [{"val": "foo"}])
+        self.db.execute("ROLLBACK")
+        self.assertEqual(self.db.execute("SELECT val FROM cs50"), [])
+
     def tearDown(self):
         self.db.execute("DROP TABLE cs50")
         self.db.execute("DROP TABLE IF EXISTS foo")
@@ -146,14 +168,16 @@ class MySQLTests(SQLTests):
     @classmethod
     def setUpClass(self):
         self.db = SQL("mysql://root@localhost/test")
+        print("\nMySQL tests")
 
     def setUp(self):
         self.db.execute("CREATE TABLE cs50 (id INTEGER NOT NULL AUTO_INCREMENT, val VARCHAR(16), bin BLOB, PRIMARY KEY (id))")
 
 class PostgresTests(SQLTests):
     @classmethod
     def setUpClass(self):
-        self.db = SQL("postgresql://postgres@localhost/test")
+        self.db = SQL("postgresql://root:test@localhost/test")
+        print("\nPOSTGRES tests")
 
     def setUp(self):
         self.db.execute("CREATE TABLE cs50 (id SERIAL PRIMARY KEY, val VARCHAR(16), bin BYTEA)")
@@ -166,6 +190,7 @@ class SQLiteTests(SQLTests):
     def setUpClass(self):
         open("test.db", "w").close()
         self.db = SQL("sqlite:///test.db")
+        print("\nSQLite tests")
 
     def setUp(self):
         self.db.execute("CREATE TABLE cs50(id INTEGER PRIMARY KEY, val TEXT, bin BLOB)")