handling connect on engine level

Kareem Zidane · Kareem Zidane · commit 87087747822a · 2018-04-26T13:30:16.000-04:00
diff --git a/.travis.yml b/.travis.yml
@@ -14,7 +14,7 @@ install:
 before_script:
 - mysql -e 'CREATE DATABASE IF NOT EXISTS test;'
 - psql -c 'create database test;' -U postgres
-- touch test.db
+- touch test.db test1.db
 script: python tests/sql.py
 after_script: rm -f test.db
 jobs:
diff --git a/src/cs50/sql.py b/src/cs50/sql.py
@@ -33,24 +33,18 @@ def __init__(self, url, **kwargs):
             if not os.path.isfile(matches.group(1)):
                 raise RuntimeError("not a file: {}".format(matches.group(1)))
 
-            # Optionally enable foreign key constraints
-            # http://docs.sqlalchemy.org/en/latest/dialects/sqlite.html#foreign-key-support
-            if kwargs.pop("pragma_foreign_keys", False):
-                @sqlalchemy.event.listens_for(sqlalchemy.engine.Engine, "connect")
-                def _set_sqlite_pragma(dbapi_connection, connection_record):
-                    """Enables foreign key support."""
+            pragma_foreign_keys = kwargs.pop("pragma_foreign_keys", False)
 
-                    # Ensure backend is sqlite
-                    if type(dbapi_connection) is sqlite3.Connection:
-                        cursor = dbapi_connection.cursor()
-
-                        # Respect foreign key constraints by default
-                        cursor.execute("PRAGMA foreign_keys=ON")
-                        cursor.close()
+            # Create engine, raising exception if back end's module not installed
+            self.engine = sqlalchemy.create_engine(url, **kwargs)
 
+            # Whether to enable foreign key constraints
+            if pragma_foreign_keys:
+                sqlalchemy.event.listen(self.engine, "connect", _on_connect)
+        else:
+            # Create engine, raising exception if back end's module not installed
+            self.engine = sqlalchemy.create_engine(url, **kwargs)
 
-        # Create engine, raising exception if back end's module not installed
-        self.engine = sqlalchemy.create_engine(url, **kwargs)
 
         # Log statements to standard error
         logging.basicConfig(level=logging.DEBUG)
@@ -229,3 +223,16 @@ def process(value):
         else:
             self.logger.debug(termcolor.colored(log, "green"))
             return ret
+
+
+# http://docs.sqlalchemy.org/en/latest/dialects/sqlite.html#foreign-key-support
+def _on_connect(dbapi_connection, connection_record):
+    """Enables foreign key support."""
+
+    # Ensure backend is sqlite
+    if type(dbapi_connection) is sqlite3.Connection:
+        cursor = dbapi_connection.cursor()
+
+        # Respect foreign key constraints by default
+        cursor.execute("PRAGMA foreign_keys=ON")
+        cursor.close()
diff --git a/tests/sql.py b/tests/sql.py
@@ -107,18 +107,26 @@ class SQLiteTests(SQLTests):
     @classmethod
     def setUpClass(self):
         self.db = SQL("sqlite:///test.db")
+        self.db1 = SQL("sqlite:///test1.db", pragma_foreign_keys=True)
 
     def setUp(self):
         self.db.execute("CREATE TABLE cs50(id INTEGER PRIMARY KEY, val TEXT)")
 
-    def multi_inserts_enabled(self):
-        return False
+    def test_foreign_key_support(self):
+        self.db.execute("CREATE TABLE foo(id INTEGER PRIMARY KEY)")
+        self.db.execute("CREATE TABLE bar(foo_id INTEGER, FOREIGN KEY (foo_id) REFERENCES foo(id))")
+        self.assertEqual(self.db.execute("INSERT INTO bar VALUES(50)"), 1)
+
+        self.db1.execute("CREATE TABLE foo(id INTEGER PRIMARY KEY)")
+        self.db1.execute("CREATE TABLE bar(foo_id INTEGER, FOREIGN KEY (foo_id) REFERENCES foo(id))")
+        self.assertEqual(self.db1.execute("INSERT INTO bar VALUES(50)"), None)
 
 if __name__ == "__main__":
     suite = unittest.TestSuite([
         unittest.TestLoader().loadTestsFromTestCase(SQLiteTests),
         unittest.TestLoader().loadTestsFromTestCase(MySQLTests),
         unittest.TestLoader().loadTestsFromTestCase(PostgresTests)
     ])
-    logging.getLogger("cs50.sql").disabled = True
+
+    logging.getLogger("cs50").disabled = True
     sys.exit(not unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful())
