[sqllab] Help sqllab forget query history

superset/assets/src/SqlLab/components/QueryAutoRefresh.jsx

@@ -8,6 +8,7 @@ const $ = require('jquery');
 
 const QUERY_UPDATE_FREQ = 2000;
 const QUERY_UPDATE_BUFFER_MS = 5000;
+const MAX_QUERY_AGE_TO_POLL = 21600000;
 
 class QueryAutoRefresh extends React.PureComponent {
   componentWillMount() {
@@ -19,10 +20,12 @@ class QueryAutoRefresh extends React.PureComponent {
   shouldCheckForQueries() {
     // if there are started or running queries, this method should return true
     const { queries } = this.props;
-    const queryKeys = Object.keys(queries);
-    const queriesAsArray = queryKeys.map(key => queries[key]);
-    return queriesAsArray.some(q =>
-      ['running', 'started', 'pending', 'fetching'].indexOf(q.state) >= 0);
+    const now = new Date().getTime();
+    return Object.values(queries)
+      .some(
+        q => ['running', 'started', 'pending', 'fetching'].indexOf(q.state) >= 0 &&
+        now - q.startDttm < MAX_QUERY_AGE_TO_POLL,
+      );
   }
   startTimer() {
     if (!(this.timer)) {

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,
       };

superset/db_engine_specs.py

@@ -652,7 +652,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
 

superset/views/core.py

@@ -25,7 +25,7 @@
 import pandas as pd
 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
@@ -2541,6 +2541,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))