From be6d1d41d3c3c2981829170210a94016a0414b46 Mon Sep 17 00:00:00 2001
From: Grace Guo <grace.guo@airbnb.com>
Date: Wed, 16 Jun 2021 16:57:31 -0700
Subject: [PATCH] fix: SQL Lab show "Refetch Results" button while fetching new
 query results (#15109)

* fix: SQL Lab show "Refetch Results" button while fetching new query results

* fix comments
---
 superset-frontend/src/SqlLab/reducers/sqlLab.js | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/superset-frontend/src/SqlLab/reducers/sqlLab.js b/superset-frontend/src/SqlLab/reducers/sqlLab.js
index daa06a97c88d..ec8438140ac3 100644
--- a/superset-frontend/src/SqlLab/reducers/sqlLab.js
+++ b/superset-frontend/src/SqlLab/reducers/sqlLab.js
@@ -517,7 +517,20 @@ export default function sqlLabReducer(state = {}, action) {
           if (changedQuery.changedOn > queriesLastUpdate) {
             queriesLastUpdate = changedQuery.changedOn;
           }
-          newQueries[id] = { ...state.queries[id], ...changedQuery };
+          const prevState = state.queries[id].state;
+          const currentState = changedQuery.state;
+          newQueries[id] = {
+            ...state.queries[id],
+            ...changedQuery,
+            // race condition:
+            // because of async behavior, sql lab may still poll a couple of seconds
+            // when it started fetching or finished rendering results
+            state:
+              currentState === 'success' &&
+              ['fetching', 'success'].includes(prevState)
+                ? prevState
+                : currentState,
+          };
           change = true;
         }
       });