From e5223ecb03f3e48f2c9d0b43a851ed9eefb74ff6 Mon Sep 17 00:00:00 2001
From: Peter Kosztolanyi <peter.kosztolanyi@gmail.com>
Date: Sun, 11 Jul 2021 01:40:16 +0100
Subject: [PATCH] Add cancel_query_on_windows_unload option to database

---
 .../src/SqlLab/components/SqlEditor.jsx        |  5 ++++-
 .../src/SqlLab/reducers/sqlLab.js              |  5 ++++-
 .../database/DatabaseModal/ExtraOptions.tsx    | 18 ++++++++++++++++++
 .../src/views/CRUD/data/database/types.ts      |  1 +
 superset/databases/api.py                      |  1 +
 tests/integration_tests/databases/api_tests.py |  1 +
 6 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/superset-frontend/src/SqlLab/components/SqlEditor.jsx b/superset-frontend/src/SqlLab/components/SqlEditor.jsx
index f869e6bb1f5f7..540cfb77f3f41 100644
--- a/superset-frontend/src/SqlLab/components/SqlEditor.jsx
+++ b/superset-frontend/src/SqlLab/components/SqlEditor.jsx
@@ -246,7 +246,10 @@ class SqlEditor extends React.PureComponent {
   }
 
   onBeforeUnload(event) {
-    if (this.props.latestQuery?.state === 'running') {
+    if (
+      this.props.database?.extra_json?.cancel_query_on_windows_unload &&
+      this.props.latestQuery?.state === 'running'
+    ) {
       event.preventDefault();
       this.stopQuery();
     }
diff --git a/superset-frontend/src/SqlLab/reducers/sqlLab.js b/superset-frontend/src/SqlLab/reducers/sqlLab.js
index daa06a97c88da..9e423ba7c65b2 100644
--- a/superset-frontend/src/SqlLab/reducers/sqlLab.js
+++ b/superset-frontend/src/SqlLab/reducers/sqlLab.js
@@ -499,7 +499,10 @@ export default function sqlLabReducer(state = {}, action) {
     [actions.SET_DATABASES]() {
       const databases = {};
       action.databases.forEach(db => {
-        databases[db.id] = db;
+        databases[db.id] = {
+          ...db,
+          extra_json: JSON.parse(db.extra || ''),
+        };
       });
       return { ...state, databases };
     },
diff --git a/superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx b/superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx
index a2a49b2f707b9..80b4044202202 100644
--- a/superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx
+++ b/superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx
@@ -294,6 +294,24 @@ const ExtraOptions = ({
             />
           </div>
         </StyledInputContainer>
+        <StyledInputContainer css={{ no_margin_bottom }}>
+          <div className="input-container">
+            <IndeterminateCheckbox
+              id="cancel_query_on_windows_unload"
+              indeterminate={false}
+              checked={!!db?.extra_json?.cancel_query_on_windows_unload}
+              onChange={onExtraInputChange}
+              labelText={t('Cancel query on window unload event')}
+            />
+            <InfoTooltip
+              tooltip={t(
+                'Terminate running queries when browser window closed or navigated ' +
+                  'to another page. Available for Presto, Hive, MySQL, Postgres and ' +
+                  'Snowflake databases.',
+              )}
+            />
+          </div>
+        </StyledInputContainer>
       </Collapse.Panel>
       <Collapse.Panel
         header={
diff --git a/superset-frontend/src/views/CRUD/data/database/types.ts b/superset-frontend/src/views/CRUD/data/database/types.ts
index 3c6fa8a825875..d582678cb4954 100644
--- a/superset-frontend/src/views/CRUD/data/database/types.ts
+++ b/superset-frontend/src/views/CRUD/data/database/types.ts
@@ -72,6 +72,7 @@ export type DatabaseObject = {
     }; // No field, holds schema and table timeout
     allows_virtual_table_explore?: boolean; // in SQL Lab
     schemas_allowed_for_csv_upload?: [] | string; // in Security
+    cancel_query_on_windows_unload?: boolean; // in Performance
     version?: string;
 
     // todo: ask beto where this should live
diff --git a/superset/databases/api.py b/superset/databases/api.py
index 8e6219fa7988e..271cd6e484404 100644
--- a/superset/databases/api.py
+++ b/superset/databases/api.py
@@ -136,6 +136,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
         "database_name",
         "explore_database_id",
         "expose_in_sqllab",
+        "extra",
         "force_ctas_schema",
         "id",
     ]
diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py
index 2c9fbb9f9cc1f..d0d3d94e04130 100644
--- a/tests/integration_tests/databases/api_tests.py
+++ b/tests/integration_tests/databases/api_tests.py
@@ -177,6 +177,7 @@ def test_get_items(self):
             "database_name",
             "explore_database_id",
             "expose_in_sqllab",
+            "extra",
             "force_ctas_schema",
             "id",
         ]