Style fixes. Minor design improvements, including removing SQL class URL variable, and always committing session so as to release locks.

Author: jsarchibald

src/cs50/sql.py

@@ -43,7 +43,7 @@ def __init__(self, url, **kwargs):
         import os
         import re
         import sqlalchemy
-        import sqlalchemy.orm as orm
+        import sqlalchemy.orm
         import sqlite3
 
         # Get logger
@@ -57,14 +57,11 @@ 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 = sqlalchemy.orm.session.sessionmaker(bind=self._engine)
         self._session = None
 
         # Listener for connections
@@ -101,6 +98,7 @@ def __del__(self):
         """Close database session and connection."""
         if self._session is not None:
             self._session.close()
+            self._session = None
 
     @_enable_logging
     def execute(self, sql, *args, **kwargs):
@@ -134,11 +132,11 @@ def execute(self, sql, *args, **kwargs):
                 command = token.value.upper()
                 break
 
-            # Begin a new transaction session, if done manually
+            # Begin a new session, if transaction started by caller (not using autocommit)
             elif token.value.upper() in ["BEGIN", "START"]:
                 if self._session is not None:
                     self._session.close()
-                self._session = self.Session()
+                self._session = self._Session()
         else:
             command = None
 
@@ -288,7 +286,7 @@ def execute(self, sql, *args, **kwargs):
         # Connect to database (for transactions' sake)
         session = self._session
         if session is None:
-            session = self.Session()
+            session = self._Session()
 
         # Set up a Flask app teardown function to close session at teardown
         try:
@@ -303,11 +301,12 @@ def execute(self, sql, *args, **kwargs):
             if not hasattr(self, "teardown_appcontext_added"):
                 self.teardown_appcontext_added = True
 
-                # Register shutdown_session on app context teardown
                 @flask.current_app.teardown_appcontext
                 def shutdown_session(exception=None):
+                    """Close any existing session on app context teardown."""
                     if self._session is not None:
                         self._session.close()
+                        self._session = None
 
         except (ModuleNotFoundError, AssertionError):
             pass
@@ -370,11 +369,9 @@ def shutdown_session(exception=None):
                     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.commit()
                     session.close()
 
             # If constraint violated, return None

tests/flask/application.py

@@ -2,19 +2,21 @@
 import os
 import requests
 import sys
-from flask import Flask, render_template
 
 sys.path.insert(0, "../../src")
 
 import cs50
 import cs50.flask
 
+from flask import Flask, render_template
+
 app = Flask(__name__)
 
 logging.disable(logging.CRITICAL)
 os.environ["WERKZEUG_RUN_MAIN"] = "true"
 
-db = cs50.SQL("sqlite:///../test.db")
+db_url = "sqlite:///../test.db"
+db = cs50.SQL(db_url)
 
 @app.route("/")
 def index():
@@ -28,7 +30,7 @@ def f():
 @app.route("/autocommit")
 def autocommit():
     db.execute("INSERT INTO test (val) VALUES (?)", "def")
-    db2 = cs50.SQL(db.url)
+    db2 = cs50.SQL(db_url)
     ret = db2.execute("SELECT val FROM test WHERE val=?", "def")
     return str(ret == [{"val": "def"}])
 
@@ -55,9 +57,9 @@ def insert():
 @app.route("/multiple_connections")
 def multiple_connections():
     ctx = len(app.teardown_appcontext_funcs)
-    db1 = cs50.SQL(db.url)
+    db1 = cs50.SQL(db_url)
     td1 = (len(app.teardown_appcontext_funcs) == ctx + 1)
-    db2 = cs50.SQL(db.url)
+    db2 = cs50.SQL(db_url)
     td2 = (len(app.teardown_appcontext_funcs) == ctx + 2)
     return str(td1 and td2)
 

tests/flask/test.py

@@ -1,11 +1,11 @@
-from application import app
 import logging
 import requests
 import sys
 import threading
 import time
 import unittest
 
+from application import app
 
 def request(route):
     r = requests.get("http://localhost:5000/{}".format(route))

tests/sql.py

@@ -120,7 +120,7 @@ def test_autocommit(self):
         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)
+        db2 = SQL(self.db_url)
         self.assertEqual(db2.execute("DELETE FROM cs50 WHERE id < 3"), 2)
 
     def test_commit(self):
@@ -129,7 +129,7 @@ def test_commit(self):
         self.db.execute("COMMIT")
 
         # Load a new database instance to confirm the INSERT was committed
-        db2 = SQL(self.db.url)
+        db2 = SQL(self.db_url)
         self.assertEqual(db2.execute("SELECT val FROM cs50"), [{"val": "foo"}])
 
     def test_rollback(self):
@@ -167,7 +167,8 @@ def tearDownClass(self):
 class MySQLTests(SQLTests):
     @classmethod
     def setUpClass(self):
-        self.db = SQL("mysql://root@localhost/test")
+        self.db_url = "mysql://root@localhost/test"
+        self.db = SQL(self.db_url)
         print("\nMySQL tests")
 
     def setUp(self):
@@ -176,7 +177,8 @@ def setUp(self):
 class PostgresTests(SQLTests):
     @classmethod
     def setUpClass(self):
-        self.db = SQL("postgresql://root:test@localhost/test")
+        self.db_url = "postgresql://root:test@localhost/test"
+        self.db = SQL(self.db_url)
         print("\nPOSTGRES tests")
 
     def setUp(self):
@@ -189,7 +191,8 @@ class SQLiteTests(SQLTests):
     @classmethod
     def setUpClass(self):
         open("test.db", "w").close()
-        self.db = SQL("sqlite:///test.db")
+        self.db_url = "sqlite:///test.db"
+        self.db = SQL(self.db_url)
         print("\nSQLite tests")
 
     def setUp(self):