From 4a26ee18139dcef875d911f9f89141393a7cfd3e Mon Sep 17 00:00:00 2001
From: Alexey Bezhan <al@fueledbylemons.com>
Date: Wed, 9 Jan 2019 12:22:51 +0000
Subject: [PATCH] Set statement timeout on all DB connections

A recent issue with a long-running query (#2288) highlighted the
fact that even though the original HTTP connection might be closed
(for example after gorouter timeout of 15 minutes, which returns a
504 response to the client), the request worker will not be stopped.

This means that the worker is spending time and potentially DB
resources generating a response that will never be delivered.

Gunicorn's timeout setting only applies to sync workers and there
doesn't seem to be an option to interrupt individual requests in
gevent/eventlet deployments.

Since the most likely (and potentially most dangerous) scenario for
this is a long-running DB query, we can set a statement timeout on
our DB connections. This will raise a sqlalchemy.exc.OperationalError
(wrapping psycopg2.extensions.QueryCanceledError), interrupting the
request after the given timeout has been reached.

This is a Postgres client setting, so the database itself will abort
the transaction when it reaches the set timeout.

Since this will also apply to our celery tasks (including potentially
long-running nightly tasks) we set a timeout of 20 minutes to begin
with.

This can potentially be split in the future to set a different value
for each app, so that we could limit API requests even more.
---
 app/__init__.py | 15 ++++++++++++++-
 app/config.py   |  1 +
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/app/__init__.py b/app/__init__.py
index b4a470ac0..1f94c697e 100644
--- a/app/__init__.py
+++ b/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()
diff --git a/app/config.py b/app/config.py
index 138d18b65..a0569534e 100644
--- a/app/config.py
+++ b/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'