fix: Improve registration of the event listener that enables FKs in sqlite

Author: natkam

grader_service/main.py

@@ -11,13 +11,14 @@
 import secrets
 import shutil
 import signal
+import sqlite3
 import subprocess
 import sys
+from typing import Any, Callable
 
 import tornado
-import uvloop as uvloop
 from jupyterhub.log import log_request
-from sqlalchemy import create_engine
+from sqlalchemy import Engine, create_engine, event
 from sqlalchemy.orm import scoped_session, sessionmaker
 from tornado.httpserver import HTTPServer
 from traitlets import (
@@ -63,6 +64,47 @@ def get_session_maker(url) -> scoped_session:
     return scoped_session(sessionmaker(bind=engine))
 
 
+def enable_foreign_keys_for_sqlite() -> Callable[[sqlite3.Connection, Any], None] | None:
+    """
+    Register a listener that enables foreign key constraints for SQLite databases.
+
+    It returns the function that sets the foreign keys pragma, and which is
+    registered with the event listener on Engine.connect. The function can be
+    used to later remove the listener:
+
+    e.g.::
+
+        from sqlalchemy import Engine, event
+
+        event.remove(Engine, "connect", set_sqlite_pragma)
+
+    This function is a no-op if the DATABASE_TYPE environment variable is not set to "sqlite".
+    In this case it returns `None`.
+    """
+    database_type = os.getenv("DATABASE_TYPE")
+    if database_type == "sqlite":
+        # The following function is needed to enable foreign key constraints in SQLite.
+        # The code was copied from the SQLAlchemy documentation:
+        # https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#foreign-key-support
+        @event.listens_for(Engine, "connect")
+        def set_sqlite_pragma(dbapi_connection, connection_record):
+            # the sqlite3 driver will not set PRAGMA foreign_keys
+            # if autocommit=False; set to True temporarily
+            ac = dbapi_connection.autocommit
+            dbapi_connection.autocommit = True
+
+            cursor = dbapi_connection.cursor()
+            # Note: this is a SQLite-specific pragma
+            cursor.execute("PRAGMA foreign_keys=ON")
+            cursor.close()
+
+            # restore previous autocommit setting
+            dbapi_connection.autocommit = ac
+
+        return set_sqlite_pragma
+    return None
+
+
 class GraderService(config.Application):
     name = "grader-service"
     version = __version__
@@ -111,13 +153,13 @@ def _default_db_url(self):
         allow_none=False,
     ).tag(config=True)
 
-    max_body_size = Int(104857600, help="Sets the max buffer size in bytes, default to 100mb").tag(
+    max_body_size = Int(104857600, help="Sets the max body size in bytes, default to 100mb").tag(
         config=True
     )
 
-    max_buffer_size = Int(104857600, help="Sets the max body size in bytes, default to 100mb").tag(
-        config=True
-    )
+    max_buffer_size = Int(
+        104857600, help="Sets the max buffer size in bytes, default to 100mb"
+    ).tag(config=True)
 
     service_git_username = Unicode("grader-service", allow_none=False).tag(config=True)
 
@@ -329,6 +371,8 @@ def initialize(self, argv, *args, **kwargs):
         self.load_config_file(self.config_file)
         self.setup_loggers(self.log_level)
 
+        enable_foreign_keys_for_sqlite()
+
         self.session_maker = get_session_maker(self.db_url)
         self.init_roles()
         # use uvloop instead of default asyncio loop
@@ -398,7 +442,10 @@ def init_roles(self):
 
                         user = db.query(User).filter(User.name == username).one_or_none()
                         if user is None:
-                            self.log.info(f"Adding new user with username {username} and display name {display_name}")
+                            self.log.info(
+                                f"Adding new user with username {username} "
+                                f"and display name {display_name}"
+                            )
                             user = User()
                             user.name = username
                             user.display_name = display_name

grader_service/orm/base.py

@@ -5,36 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 import enum
-import os
 
-from sqlalchemy import event
-from sqlalchemy.engine import Engine
 from sqlalchemy.orm import declarative_base
 
 from grader_service.api.models.base_model import Model
 
 Base = declarative_base()
 
-database_type = os.getenv("DATABASE_TYPE")
-if database_type == "sqlite":
-    # The following function is needed to enable foreign key constraints in SQLite;
-    # The code was copied from the SQLAlchemy documentation:
-    # https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#foreign-key-support
-    @event.listens_for(Engine, "connect")
-    def set_sqlite_pragma(dbapi_connection, connection_record):
-        # the sqlite3 driver will not set PRAGMA foreign_keys
-        # if autocommit=False; set to True temporarily
-        ac = dbapi_connection.autocommit
-        dbapi_connection.autocommit = True
-
-        cursor = dbapi_connection.cursor()
-        # Note: this is a SQLite-specific pragma
-        cursor.execute("PRAGMA foreign_keys=ON")
-        cursor.close()
-
-        # restore previous autocommit setting
-        dbapi_connection.autocommit = ac
-
 
 class Serializable(object):
     @property

grader_service/tests/conftest.py

@@ -10,13 +10,12 @@
 import pytest
 from alembic import config
 from alembic.command import upgrade
-from sqlalchemy import event
-from sqlalchemy.engine import Engine
+from sqlalchemy import Engine, event
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 
 from grader_service import GraderService, handlers
 from grader_service.auth.dummy import DummyAuthenticator
-from grader_service.main import get_session_maker
+from grader_service.main import enable_foreign_keys_for_sqlite, get_session_maker
 from grader_service.orm import User
 from grader_service.registry import HandlerPathRegistry
 from grader_service.server import GraderServer
@@ -34,37 +33,22 @@ def set_database_type_to_sqlite():
     which enables foreign keys support for SQLite database.
     Unset the var after the test runs.
 
-    This is a hack, because normally `DATABASE_TYPE` is not set when tests run.
+    This is a necessary hack, because normally `DATABASE_TYPE` is not set when tests run.
     It also cannot be set to "sqlite" by default, because we have some tests
     running on PostgreSQL, where executing the sqlite pragma would cause an error.
 
     Outside the test setting, the DATABASE_TYPE environment variable is set, and
-    the event listener is registered in grader_service/orm/base.py.
+    the event listener is registered in `grader_service/orm/base.py`.
     """
     os.environ["DATABASE_TYPE"] = "sqlite"
-
-    # Original code taken from the SQLAlchemy documentation, here slightly adjusted:
-    # https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#foreign-key-support
-    @event.listens_for(Engine, "connect")
-    def set_sqlite_pragma(dbapi_connection, connection_record):
-        database_type = os.getenv("DATABASE_TYPE")
-        if database_type == "sqlite":
-            # the sqlite3 driver will not set PRAGMA foreign_keys
-            # if autocommit=False; set to True temporarily
-            ac = dbapi_connection.autocommit
-            dbapi_connection.autocommit = True
-
-            cursor = dbapi_connection.cursor()
-            # Note: this is a SQLite-specific pragma
-            cursor.execute("PRAGMA foreign_keys=ON")
-            cursor.close()
-
-            # restore previous autocommit setting
-            dbapi_connection.autocommit = ac
-
+    set_sqlite_pragma = enable_foreign_keys_for_sqlite()
+    assert event.contains(Engine, "connect", set_sqlite_pragma)
     yield
+
     # Unset the variable for other tests, which may use a different database type
     os.environ.pop("DATABASE_TYPE")
+    # Remove the event listener to avoid interference with other tests
+    event.remove(Engine, "connect", set_sqlite_pragma)
 
 
 @pytest.fixture(scope="function")

grader_service/tests/handlers/test_assignment_handler.py

@@ -33,8 +33,6 @@ def test_foreign_key_constraints_in_sqlite(
     engine = session.get_bind()
     sub = insert_submission(engine, a_id, "ubuntu", 1)
 
-    session = sql_alchemy_sessionmaker()
-
     with pytest.raises(IntegrityError, match="FOREIGN KEY constraint failed"):
         # Try to delete an existing assignment with a submission
         session.query(AssignmentORM).filter(AssignmentORM.id == a_id).delete()