From 48100172ef54a45fe4754742d4b864476ba10737 Mon Sep 17 00:00:00 2001
From: Timi Fasubaa <timi.fasubaa@airbnb.com>
Date: Fri, 20 Apr 2018 11:21:08 -0700
Subject: [PATCH] make queries older than 6 hours timeout

---
 .../SqlLab/components/QueryAutoRefresh.jsx    | 14 +++++----
 superset/assets/src/SqlLab/reducers.js        |  5 +--
 superset/db_engine_specs.py                   |  2 +-
 superset/views/core.py                        | 31 ++++++++++++++++++-
 4 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/superset/assets/src/SqlLab/components/QueryAutoRefresh.jsx b/superset/assets/src/SqlLab/components/QueryAutoRefresh.jsx
index e026575fbda7b..55e06cc1467b2 100644
--- a/superset/assets/src/SqlLab/components/QueryAutoRefresh.jsx
+++ b/superset/assets/src/SqlLab/components/QueryAutoRefresh.jsx
@@ -8,7 +8,7 @@ const $ = require('jquery');
 
 const QUERY_UPDATE_FREQ = 2000;
 const QUERY_UPDATE_BUFFER_MS = 5000;
-const QUERY_POLL_WINDOW = 21600000; // 6 hours.
+const MAX_QUERY_AGE_TO_POLL = 21600000;
 
 class QueryAutoRefresh extends React.PureComponent {
   componentWillMount() {
@@ -18,12 +18,14 @@ class QueryAutoRefresh extends React.PureComponent {
     this.stopTimer();
   }
   shouldCheckForQueries() {
-    // if there are started or running queries < 6 hours old, this method should return true
+    // if there are started or running queries, this method should return true
     const { queries } = this.props;
+    const now = new Date().getTime();
     return Object.values(queries)
-      .filter(q => (q.startDttm >= this.props.queriesLastUpdate - QUERY_POLL_WINDOW))
-      .some(q =>
-      ['running', 'started', 'pending', 'fetching'].indexOf(q.state) >= 0);
+      .some(
+        q => ['running', 'started', 'pending', 'fetching'].indexOf(q.state) >= 0 &&
+        now - q.startDttm < MAX_QUERY_AGE_TO_POLL,
+      );
   }
   startTimer() {
     if (!(this.timer)) {
@@ -35,7 +37,7 @@ class QueryAutoRefresh extends React.PureComponent {
     this.timer = null;
   }
   stopwatch() {
-    // only poll /superset/queries/ if there are started or running queries started <6 hours ago
+    // only poll /superset/queries/ if there are started or running queries
     if (this.shouldCheckForQueries()) {
       const url = `/superset/queries/${this.props.queriesLastUpdate - QUERY_UPDATE_BUFFER_MS}`;
       $.getJSON(url, (data) => {
diff --git a/superset/assets/src/SqlLab/reducers.js b/superset/assets/src/SqlLab/reducers.js
index f01f2c3bb73a2..9c67f3822cf3f 100644
--- a/superset/assets/src/SqlLab/reducers.js
+++ b/superset/assets/src/SqlLab/reducers.js
@@ -162,9 +162,6 @@ export const sqlLabReducer = function (state, action) {
       return alterInObject(state, 'queries', action.query, { state: 'fetching' });
     },
     [actions.QUERY_SUCCESS]() {
-      if (action.query.state === 'stopped') {
-        return state;
-      }
       let rows;
       if (action.results.data) {
         rows = action.results.data.length;
@@ -174,7 +171,7 @@ export const sqlLabReducer = function (state, action) {
         progress: 100,
         results: action.results,
         rows,
-        state: 'success',
+        state: action.query.state,
         errorMessage: null,
         cached: false,
       };
diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py
index 9d2c65db9fae0..a718a0d62ca85 100644
--- a/superset/db_engine_specs.py
+++ b/superset/db_engine_specs.py
@@ -675,7 +675,7 @@ def handle_cursor(cls, cursor, query, session):
             stats = polled.get('stats', {})
 
             query = session.query(type(query)).filter_by(id=query.id).one()
-            if query.status == QueryStatus.STOPPED:
+            if query.status in [QueryStatus.STOPPED, QueryStatus.TIMED_OUT]:
                 cursor.cancel()
                 break
 
diff --git a/superset/views/core.py b/superset/views/core.py
index 0359cae739118..d290ecc9ca16f 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -26,7 +26,7 @@
 import simplejson as json
 from six import text_type
 import sqlalchemy as sqla
-from sqlalchemy import create_engine
+from sqlalchemy import and_, create_engine, update
 from sqlalchemy.engine.url import make_url
 from sqlalchemy.exc import IntegrityError
 from unidecode import unidecode
@@ -2549,6 +2549,35 @@ def queries(self, last_updated_ms):
             .all()
         )
         dict_queries = {q.client_id: q.to_dict() for q in sql_queries}
+
+        now = int(round(time.time() * 1000))
+
+        unfinished_states = [
+            utils.QueryStatus.PENDING,
+            utils.QueryStatus.RUNNING,
+        ]
+
+        queries_to_timeout = [
+            client_id for client_id, query_dict in dict_queries.items()
+            if (
+                query_dict['state'] in unfinished_states and (
+                    now - query_dict['startDttm'] >
+                    config.get('SQLLAB_ASYNC_TIME_LIMIT_SEC') * 1000
+                )
+            )
+        ]
+
+        if queries_to_timeout:
+            update(Query).where(
+                and_(
+                    Query.user_id == g.user.get_id(),
+                    Query.client_id in queries_to_timeout,
+                ),
+            ).values(state=utils.QueryStatus.TIMED_OUT)
+
+            for client_id in queries_to_timeout:
+                dict_queries[client_id]['status'] = utils.QueryStatus.TIMED_OUT
+
         return json_success(
             json.dumps(dict_queries, default=utils.json_int_dttm_ser))