From 9ae81b7c33d63873fdf2c4ff5c579b15a934ad9b Mon Sep 17 00:00:00 2001
From: "JUST.in DO IT" <justin.park@airbnb.com>
Date: Mon, 13 Mar 2023 16:03:19 -0700
Subject: [PATCH] fix(sqllab): empty large query results from localStorage
 (#23302)

---
 superset-frontend/src/SqlLab/constants.ts     |  1 +
 .../SqlLab/utils/emptyQueryResults.test.js    | 39 ++++++++++++++++++-
 .../utils/reduxStateToLocalStorageHelper.js   | 21 ++++++++--
 3 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/superset-frontend/src/SqlLab/constants.ts b/superset-frontend/src/SqlLab/constants.ts
index 108ce8956..03f725c09 100644
--- a/superset-frontend/src/SqlLab/constants.ts
+++ b/superset-frontend/src/SqlLab/constants.ts
@@ -83,6 +83,7 @@ export const BYTES_PER_CHAR = 2;
 // browser's localStorage max usage constants
 export const LOCALSTORAGE_MAX_QUERY_AGE_MS = 24 * 60 * 60 * 1000; // 24 hours
 export const LOCALSTORAGE_MAX_USAGE_KB = 5 * 1024; // 5M
+export const LOCALSTORAGE_MAX_QUERY_RESULTS_KB = 1 * 1024; // 1M
 export const LOCALSTORAGE_WARNING_THRESHOLD = 0.9;
 export const LOCALSTORAGE_WARNING_MESSAGE_THROTTLE_MS = 8000; // danger type toast duration
 
diff --git a/superset-frontend/src/SqlLab/utils/emptyQueryResults.test.js b/superset-frontend/src/SqlLab/utils/emptyQueryResults.test.js
index 6a441062b..78072f3e9 100644
--- a/superset-frontend/src/SqlLab/utils/emptyQueryResults.test.js
+++ b/superset-frontend/src/SqlLab/utils/emptyQueryResults.test.js
@@ -20,7 +20,12 @@ import {
   emptyQueryResults,
   clearQueryEditors,
 } from 'src/SqlLab/utils/reduxStateToLocalStorageHelper';
-import { LOCALSTORAGE_MAX_QUERY_AGE_MS } from 'src/SqlLab/constants';
+import {
+  KB_STORAGE,
+  BYTES_PER_CHAR,
+  LOCALSTORAGE_MAX_QUERY_AGE_MS,
+  LOCALSTORAGE_MAX_QUERY_RESULTS_KB,
+} from 'src/SqlLab/constants';
 import { queries, defaultQueryEditor } from '../fixtures';
 
 describe('reduxStateToLocalStorageHelper', () => {
@@ -45,6 +50,38 @@ describe('reduxStateToLocalStorageHelper', () => {
     expect(emptiedQuery[id].results).toEqual({});
   });
 
+  it('should empty query.results if query,.results size is greater than LOCALSTORAGE_MAX_QUERY_RESULTS_KB', () => {
+    const reasonableSizeQuery = {
+      ...queries[0],
+      startDttm: Date.now(),
+      results: { data: [{ a: 1 }] },
+    };
+    const largeQuery = {
+      ...queries[1],
+      startDttm: Date.now(),
+      results: {
+        data: [
+          {
+            jsonValue: `{"str":"${new Array(
+              (LOCALSTORAGE_MAX_QUERY_RESULTS_KB / BYTES_PER_CHAR) * KB_STORAGE,
+            )
+              .fill(0)
+              .join('')}"}`,
+          },
+        ],
+      },
+    };
+    expect(Object.keys(largeQuery.results)).toContain('data');
+    const emptiedQuery = emptyQueryResults({
+      [reasonableSizeQuery.id]: reasonableSizeQuery,
+      [largeQuery.id]: largeQuery,
+    });
+    expect(emptiedQuery[largeQuery.id].results).toEqual({});
+    expect(emptiedQuery[reasonableSizeQuery.id].results).toEqual(
+      reasonableSizeQuery.results,
+    );
+  });
+
   it('should only return selected keys for query editor', () => {
     const queryEditors = [defaultQueryEditor];
     expect(Object.keys(queryEditors[0])).toContain('schemaOptions');
diff --git a/superset-frontend/src/SqlLab/utils/reduxStateToLocalStorageHelper.js b/superset-frontend/src/SqlLab/utils/reduxStateToLocalStorageHelper.js
index 66e33b07c..d03354f16 100644
--- a/superset-frontend/src/SqlLab/utils/reduxStateToLocalStorageHelper.js
+++ b/superset-frontend/src/SqlLab/utils/reduxStateToLocalStorageHelper.js
@@ -16,7 +16,12 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { LOCALSTORAGE_MAX_QUERY_AGE_MS } from '../constants';
+import {
+  BYTES_PER_CHAR,
+  KB_STORAGE,
+  LOCALSTORAGE_MAX_QUERY_AGE_MS,
+  LOCALSTORAGE_MAX_QUERY_RESULTS_KB,
+} from '../constants';
 
 const PERSISTENT_QUERY_EDITOR_KEYS = new Set([
   'remoteId',
@@ -36,13 +41,21 @@ const PERSISTENT_QUERY_EDITOR_KEYS = new Set([
   'hideLeftBar',
 ]);
 
+function shouldEmptyQueryResults(query) {
+  const { startDttm, results } = query;
+  return (
+    Date.now() - startDttm > LOCALSTORAGE_MAX_QUERY_AGE_MS ||
+    ((JSON.stringify(results)?.length || 0) * BYTES_PER_CHAR) / KB_STORAGE >
+      LOCALSTORAGE_MAX_QUERY_RESULTS_KB
+  );
+}
+
 export function emptyQueryResults(queries) {
   return Object.keys(queries).reduce((accu, key) => {
-    const { startDttm, results } = queries[key];
+    const { results } = queries[key];
     const query = {
       ...queries[key],
-      results:
-        Date.now() - startDttm > LOCALSTORAGE_MAX_QUERY_AGE_MS ? {} : results,
+      results: shouldEmptyQueryResults(queries[key]) ? {} : results,
     };
 
     const updatedQueries = {