added support for Flask and transactions

Author: dmalan

src/cs50/sql.py

@@ -56,31 +56,23 @@ def __init__(self, url, **kwargs):
             if not os.path.isfile(matches.group(1)):
                 raise RuntimeError("not a file: {}".format(matches.group(1)))
 
-            # Create engine, disabling SQLAlchemy's own autocommit mode, raising exception if back end's module not installed
-            engine = sqlalchemy.create_engine(url, **kwargs).execution_options(autocommit=False)
+        # 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):
+        # Listener for connections
+        def connect(dbapi_connection, connection_record):
 
-                # Disable underlying API's own emitting of BEGIN and COMMIT
-                dbapi_connection.isolation_level = None
+            # Disable underlying API's own emitting of BEGIN and COMMIT
+            dbapi_connection.isolation_level = None
 
-                # Enable foreign key constraints
-                if type(dbapi_connection) is sqlite3.Connection:  # If back end is sqlite
-                    cursor = dbapi_connection.cursor()
-                    cursor.execute("PRAGMA foreign_keys=ON")
-                    cursor.close()
+            # Enable foreign key constraints
+            if type(dbapi_connection) is sqlite3.Connection:  # If back end is sqlite
+                cursor = dbapi_connection.cursor()
+                cursor.execute("PRAGMA foreign_keys=ON")
+                cursor.close()
 
-            # Register listener
-            sqlalchemy.event.listen(engine, "connect", connect)
-
-        else:
-
-            # Create engine, raising exception if back end's module not installed
-            engine = sqlalchemy.create_engine(url, **kwargs)
-
-        # Connect to database (for transactions' sake)
-        self._connection = engine.connect().execution_options(autocommit=False)
+        # Register listener
+        sqlalchemy.event.listen(self._engine, "connect", connect)
 
         # Log statements to standard error
         logging.basicConfig(level=logging.DEBUG)
@@ -266,6 +258,39 @@ def execute(self, sql, *args, **kwargs):
         # Join tokens into statement
         statement = "".join([str(token) for token in tokens])
 
+        # Connect to database (for transactions' sake)
+        try:
+
+            # Infer whether Flask is installed
+            import flask
+
+            # 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
+                @flask.current_app.teardown_appcontext
+                def shutdown_session(exception=None):
+                    print("DELETING")
+                    flask.g._connection.close()
+
+            # Use this connection
+            connection = flask.g._connection
+
+        except (ModuleNotFoundError, AssertionError):
+
+            # If no connection yet
+            if not hasattr(self, "_connection"):
+                self._connection = self._engine.connect()
+
+            # Use this connection
+            connection = self._connection
+
         # Catch SQLAlchemy warnings
         with warnings.catch_warnings():
 
@@ -279,7 +304,7 @@ def execute(self, sql, *args, **kwargs):
                 _statement = "".join([str(bytes) if token.ttype == sqlparse.tokens.Other else str(token) for token in tokens])
 
                 # Execute statement
-                result = self._connection.execute(sqlalchemy.text(statement))
+                result = connection.execute(sqlalchemy.text(statement))
 
                 # Return value
                 ret = True
@@ -310,9 +335,9 @@ def execute(self, sql, *args, **kwargs):
 
                     # If INSERT, return primary key value for a newly inserted row (or None if none)
                     elif value == "INSERT":
-                        if self._connection.engine.url.get_backend_name() in ["postgres", "postgresql"]:
+                        if self._engine.url.get_backend_name() in ["postgres", "postgresql"]:
                             try:
-                                result = self._connection.execute("SELECT LASTVAL()")
+                                result = connection.execute("SELECT LASTVAL()")
                                 ret = result.first()[0]
                             except sqlalchemy.exc.OperationalError:  # If lastval is not yet defined in this session
                                 ret = None
@@ -362,13 +387,13 @@ def __escape(value):
             if type(value) is bool:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.Number,
-                    sqlalchemy.types.Boolean().literal_processor(self._connection.engine.dialect)(value))
+                    sqlalchemy.types.Boolean().literal_processor(self._engine.dialect)(value))
 
             # bytearray, bytes
             elif type(value) in [bytearray, bytes]:
-                if self._connection.engine.url.get_backend_name() in ["mysql", "sqlite"]:
+                if self._engine.url.get_backend_name() in ["mysql", "sqlite"]:
                     return sqlparse.sql.Token(sqlparse.tokens.Other, f"x'{value.hex()}'")  # https://dev.mysql.com/doc/refman/8.0/en/hexadecimal-literals.html
-                elif self._connection.engine.url.get_backend_name() == "postgresql":
+                elif self._engine.url.get_backend_name() == "postgresql":
                     return sqlparse.sql.Token(sqlparse.tokens.Other, f"'\\x{value.hex()}'")  # https://dba.stackexchange.com/a/203359
                 else:
                     raise RuntimeError("unsupported value: {}".format(value))
@@ -377,43 +402,43 @@ def __escape(value):
             elif type(value) is datetime.date:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.String,
-                    sqlalchemy.types.String().literal_processor(self._connection.engine.dialect)(value.strftime("%Y-%m-%d")))
+                    sqlalchemy.types.String().literal_processor(self._engine.dialect)(value.strftime("%Y-%m-%d")))
 
             # datetime.datetime
             elif type(value) is datetime.datetime:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.String,
-                    sqlalchemy.types.String().literal_processor(self._connection.engine.dialect)(value.strftime("%Y-%m-%d %H:%M:%S")))
+                    sqlalchemy.types.String().literal_processor(self._engine.dialect)(value.strftime("%Y-%m-%d %H:%M:%S")))
 
             # datetime.time
             elif type(value) is datetime.time:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.String,
-                    sqlalchemy.types.String().literal_processor(self._connection.engine.dialect)(value.strftime("%H:%M:%S")))
+                    sqlalchemy.types.String().literal_processor(self._engine.dialect)(value.strftime("%H:%M:%S")))
 
             # float
             elif type(value) is float:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.Number,
-                    sqlalchemy.types.Float().literal_processor(self._connection.engine.dialect)(value))
+                    sqlalchemy.types.Float().literal_processor(self._engine.dialect)(value))
 
             # int
             elif type(value) is int:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.Number,
-                    sqlalchemy.types.Integer().literal_processor(self._connection.engine.dialect)(value))
+                    sqlalchemy.types.Integer().literal_processor(self._engine.dialect)(value))
 
             # str
             elif type(value) is str:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.String,
-                    sqlalchemy.types.String().literal_processor(self._connection.engine.dialect)(value))
+                    sqlalchemy.types.String().literal_processor(self._engine.dialect)(value))
 
             # None
             elif value is None:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.Keyword,
-                    sqlalchemy.types.NullType().literal_processor(self._connection.engine.dialect)(value))
+                    sqlalchemy.types.NullType().literal_processor(self._engine.dialect)(value))
 
             # Unsupported value
             else: