Merge pull request #2291 from alphagov/set-db-statement-timeout

Set statement timeout on all DB connections

app/__init__.py

@@ -4,7 +4,7 @@ import string
 import uuid
 
 from flask import _request_ctx_stack, request, g, jsonify
-from flask_sqlalchemy import SQLAlchemy
+from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy
 from flask_marshmallow import Marshmallow
 from flask_migrate import Migrate
 from time import monotonic
@@ -27,6 +27,19 @@ from app.encryption import Encryption
 DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 DATE_FORMAT = "%Y-%m-%d"
 
+
+class SQLAlchemy(_SQLAlchemy):
+    """We need to subclass SQLAlchemy in order to override create_engine options"""
+
+    def apply_driver_hacks(self, app, info, options):
+        super().apply_driver_hacks(app, info, options)
+        if 'connect_args' not in options:
+            options['connect_args'] = {}
+        options['connect_args']["options"] = "-c statement_timeout={}".format(
+            int(app.config['SQLALCHEMY_STATEMENT_TIMEOUT']) * 1000
+        )
+
+
 db = SQLAlchemy()
 migrate = Migrate()
 ma = Marshmallow()

app/config.py

@@ -122,6 +122,7 @@ class Config(object):
     SQLALCHEMY_POOL_SIZE = int(os.environ.get('SQLALCHEMY_POOL_SIZE', 5))
     SQLALCHEMY_POOL_TIMEOUT = 30
     SQLALCHEMY_POOL_RECYCLE = 300
+    SQLALCHEMY_STATEMENT_TIMEOUT = 1200
     PAGE_SIZE = 50
     API_PAGE_SIZE = 250
     TEST_MESSAGE_FILENAME = 'Test message'