adding a timeout limit to web api queries

Author: CamDavidsonPilon

core/pioreactor/web/app.py

@@ -4,6 +4,7 @@
 from base64 import b64decode
 from datetime import datetime
 from datetime import timezone
+from time import monotonic
 
 from flask import Flask
 from flask import g
@@ -27,6 +28,8 @@
 VERSION = __version__
 HOSTNAME = get_unit_name()
 NAME = f"pioreactor-{HOSTNAME}-api"
+APP_DATABASE_QUERY_TIMEOUT_SECONDS = 30.0
+APP_DATABASE_QUERY_PROGRESS_HANDLER_INSTRUCTIONS = 50_000
 
 
 # set up logging
@@ -232,12 +235,34 @@ def query_app_db(
 ) -> dict[str, t.Any] | list[dict[str, t.Any]] | None:
     assert am_I_leader()
     con = _get_app_db_connection()
+    cur: sqlite3.Cursor | None = None
+    query_deadline = monotonic() + APP_DATABASE_QUERY_TIMEOUT_SECONDS
+    query_timed_out = False
+
+    def query_exceeded_execution_deadline() -> int:
+        nonlocal query_timed_out
+        query_timed_out = monotonic() >= query_deadline
+        return int(query_timed_out)
+
     try:
         con.execute("PRAGMA query_only = 1")
+        con.set_progress_handler(
+            query_exceeded_execution_deadline,
+            APP_DATABASE_QUERY_PROGRESS_HANDLER_INSTRUCTIONS,
+        )
         cur = con.execute(query, args)
         rv = cur.fetchall()
-        cur.close()
+    except sqlite3.OperationalError as e:
+        if query_timed_out:
+            raise TimeoutError(
+                f"Database query exceeded {APP_DATABASE_QUERY_TIMEOUT_SECONDS:g} seconds."
+            ) from e
+        raise
     finally:
+        con.set_progress_handler(None, 0)
+        if cur is not None:
+            cur.close()
+
         # Restore to default to allow mutations via modify_app_db within same request
         try:
             con.execute("PRAGMA query_only = 0")

core/tests/web/test_db_querying.py

@@ -50,6 +50,45 @@ def test_query_app_db_disallows_dml_via_query_only_pragma(app, tmp_path) -> None
     assert _count_rows(db_path) == 1
 
 
+def test_query_app_db_interrupts_queries_that_exceed_the_execution_deadline(app, monkeypatch) -> None:
+    from pioreactor.web import app as web_app
+
+    monkeypatch.setattr(web_app, "APP_DATABASE_QUERY_TIMEOUT_SECONDS", 0.0, raising=False)
+
+    with app.app_context():
+        connection = sqlite3.connect(":memory:")
+        connection.row_factory = web_app._make_dicts
+        g._app_database = connection
+
+        with pytest.raises(TimeoutError, match="Database query exceeded"):
+            web_app.query_app_db(
+                """
+                WITH RECURSIVE counter(value) AS (
+                    VALUES(0)
+                    UNION ALL
+                    SELECT value + 1 FROM counter WHERE value < 1000000
+                )
+                SELECT SUM(value) FROM counter
+                """
+            )
+
+        # The interrupted statement must release the connection and remove its
+        # progress handler so the same request can continue using SQLite.
+        result = connection.execute(
+            """
+            WITH RECURSIVE counter(value) AS (
+                VALUES(0)
+                UNION ALL
+                SELECT value + 1 FROM counter WHERE value < 100000
+            )
+            SELECT SUM(value) AS total FROM counter
+            """
+        ).fetchone()
+
+        assert result == {"total": 5000050000}
+        connection.execute("CREATE TABLE after_timeout (value INTEGER)")
+
+
 def test_modify_app_db_rolls_back_after_integrity_error(app, tmp_path) -> None:
     from pioreactor.web import app as web_app