From a44c99899a66980934ae5459c28720063fc0abab Mon Sep 17 00:00:00 2001 From: "JUST.in DO IT" Date: Fri, 11 Aug 2023 13:23:07 -0700 Subject: [PATCH] chore(sqllab): Relocate get bootstrap data logic (#24936) --- .../src/SqlLab/actions/sqlLab.js | 9 +-- .../components/TabbedSqlEditors/index.jsx | 2 +- .../src/SqlLab/reducers/getInitialState.js | 3 +- .../SqlLab/reducers/getInitialState.test.ts | 1 - superset/sqllab/utils.py | 63 +++++++++++++++++ superset/views/base.py | 1 + superset/views/core.py | 67 +------------------ tests/integration_tests/core_tests.py | 4 +- 8 files changed, 74 insertions(+), 76 deletions(-) diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js index 5d9ecdacdf..8d39d3bbdb 100644 --- a/superset-frontend/src/SqlLab/actions/sqlLab.js +++ b/superset-frontend/src/SqlLab/actions/sqlLab.js @@ -570,15 +570,10 @@ export function addQueryEditor(queryEditor) { export function addNewQueryEditor() { return function (dispatch, getState) { const { - sqlLab: { - queryEditors, - tabHistory, - unsavedQueryEditor, - defaultDbId, - databases, - }, + sqlLab: { queryEditors, tabHistory, unsavedQueryEditor, databases }, common, } = getState(); + const defaultDbId = common.conf.SQLLAB_DEFAULT_DBID; const activeQueryEditor = queryEditors.find( qe => qe.id === tabHistory[tabHistory.length - 1], ); diff --git a/superset-frontend/src/SqlLab/components/TabbedSqlEditors/index.jsx b/superset-frontend/src/SqlLab/components/TabbedSqlEditors/index.jsx index d914715630..95d0c2529b 100644 --- a/superset-frontend/src/SqlLab/components/TabbedSqlEditors/index.jsx +++ b/superset-frontend/src/SqlLab/components/TabbedSqlEditors/index.jsx @@ -340,7 +340,7 @@ function mapStateToProps({ sqlLab, common }) { queries: sqlLab.queries, tabHistory: sqlLab.tabHistory, tables: sqlLab.tables, - defaultDbId: sqlLab.defaultDbId, + defaultDbId: common.conf.SQLLAB_DEFAULT_DBID, displayLimit: common.conf.DISPLAY_MAX_ROW, offline: sqlLab.offline, defaultQueryLimit: common.conf.DEFAULT_SQLLAB_LIMIT, diff --git a/superset-frontend/src/SqlLab/reducers/getInitialState.js b/superset-frontend/src/SqlLab/reducers/getInitialState.js index edb778fcfa..1e3ac94d11 100644 --- a/superset-frontend/src/SqlLab/reducers/getInitialState.js +++ b/superset-frontend/src/SqlLab/reducers/getInitialState.js @@ -28,7 +28,6 @@ export function dedupeTabHistory(tabHistory) { } export default function getInitialState({ - defaultDbId, common, active_tab: activeTab, tab_state_ids: tabStateIds = [], @@ -55,7 +54,7 @@ export default function getInitialState({ latestQueryId: null, autorun: false, templateParams: null, - dbId: defaultDbId, + dbId: common.conf.SQLLAB_DEFAULT_DBID, queryLimit: common.conf.DEFAULT_SQLLAB_LIMIT, validationResult: { id: null, diff --git a/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts b/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts index c06633c6f5..a3c71cbd88 100644 --- a/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts +++ b/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts @@ -20,7 +20,6 @@ import getInitialState, { dedupeTabHistory } from './getInitialState'; const apiData = { - defaultDbId: 1, common: { conf: { DEFAULT_SQLLAB_LIMIT: 1, diff --git a/superset/sqllab/utils.py b/superset/sqllab/utils.py index abceaaf136..d308bd4639 100644 --- a/superset/sqllab/utils.py +++ b/superset/sqllab/utils.py @@ -14,11 +14,31 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + from typing import Any import pyarrow as pa +from superset import db, is_feature_enabled from superset.common.db_query_status import QueryStatus +from superset.daos.database import DatabaseDAO +from superset.models.sql_lab import Query, TabState + +DATABASE_KEYS = [ + "allow_file_upload", + "allow_ctas", + "allow_cvas", + "allow_dml", + "allow_run_async", + "allows_subquery", + "backend", + "database_name", + "expose_in_sqllab", + "force_ctas_schema", + "id", + "disable_data_preview", +] def apply_display_max_row_configuration_if_require( # pylint: disable=invalid-name @@ -56,3 +76,46 @@ def write_ipc_buffer(table: pa.Table) -> pa.Buffer: writer.write_table(table) return sink.getvalue() + + +def bootstrap_sqllab_data(user_id: int | None) -> dict[str, Any]: + # send list of tab state ids + tabs_state = ( + db.session.query(TabState.id, TabState.label).filter_by(user_id=user_id).all() + ) + tab_state_ids = [str(tab_state[0]) for tab_state in tabs_state] + # return first active tab, or fallback to another one if no tab is active + active_tab = ( + db.session.query(TabState) + .filter_by(user_id=user_id) + .order_by(TabState.active.desc()) + .first() + ) + + databases: dict[int, Any] = {} + for database in DatabaseDAO.find_all(): + databases[database.id] = { + k: v for k, v in database.to_json().items() if k in DATABASE_KEYS + } + databases[database.id]["backend"] = database.backend + queries: dict[str, Any] = {} + + # These are unnecessary if sqllab backend persistence is disabled + if is_feature_enabled("SQLLAB_BACKEND_PERSISTENCE"): + # return all user queries associated with existing SQL editors + user_queries = ( + db.session.query(Query) + .filter_by(user_id=user_id) + .filter(Query.sql_editor_id.in_(tab_state_ids)) + .all() + ) + queries = { + query.client_id: dict(query.to_dict().items()) for query in user_queries + } + + return { + "tab_state_ids": tabs_state, + "active_tab": active_tab.to_dict() if active_tab else None, + "databases": databases, + "queries": queries, + } diff --git a/superset/views/base.py b/superset/views/base.py index 34232a6e34..a2c62df41b 100644 --- a/superset/views/base.py +++ b/superset/views/base.py @@ -95,6 +95,7 @@ "SQL_MAX_ROW", "SUPERSET_WEBSERVER_DOMAINS", "SQLLAB_SAVE_WARNING_MESSAGE", + "SQLLAB_DEFAULT_DBID", "DISPLAY_MAX_ROW", "GLOBAL_ASYNC_QUERIES_TRANSPORT", "GLOBAL_ASYNC_QUERIES_POLLING_DELAY", diff --git a/superset/views/core.py b/superset/views/core.py index 0394c90dd3..d377f94d65 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -48,7 +48,6 @@ from superset.connectors.base.models import BaseDatasource from superset.connectors.sqla.models import SqlaTable from superset.daos.chart import ChartDAO -from superset.daos.database import DatabaseDAO from superset.daos.datasource import DatasourceDAO from superset.dashboards.commands.importers.v0 import ImportDashboardsCommand from superset.dashboards.permalink.commands.get import GetDashboardPermalinkCommand @@ -69,8 +68,9 @@ from superset.models.core import Database from superset.models.dashboard import Dashboard from superset.models.slice import Slice -from superset.models.sql_lab import Query, TabState +from superset.models.sql_lab import Query from superset.models.user_attributes import UserAttribute +from superset.sqllab.utils import bootstrap_sqllab_data from superset.superset_typing import FlaskResponse from superset.tasks.async_queries import load_explore_json_into_cache from superset.utils import core as utils @@ -115,21 +115,6 @@ stats_logger = config["STATS_LOGGER"] logger = logging.getLogger(__name__) -DATABASE_KEYS = [ - "allow_file_upload", - "allow_ctas", - "allow_cvas", - "allow_dml", - "allow_run_async", - "allows_subquery", - "backend", - "database_name", - "expose_in_sqllab", - "force_ctas_schema", - "id", - "disable_data_preview", -] - DATASOURCE_MISSING_ERR = __("The data source seems to have been deleted") USER_MISSING_ERR = __("The user seems to have been deleted") PARAMETER_MISSING_ERR = __( @@ -1019,51 +1004,6 @@ def profile(self) -> FlaskResponse: ), ) - @staticmethod - def _get_sqllab_tabs(user_id: int | None) -> dict[str, Any]: - # send list of tab state ids - tabs_state = ( - db.session.query(TabState.id, TabState.label) - .filter_by(user_id=user_id) - .all() - ) - tab_state_ids = [str(tab_state[0]) for tab_state in tabs_state] - # return first active tab, or fallback to another one if no tab is active - active_tab = ( - db.session.query(TabState) - .filter_by(user_id=user_id) - .order_by(TabState.active.desc()) - .first() - ) - - databases: dict[int, Any] = {} - for database in DatabaseDAO.find_all(): - databases[database.id] = { - k: v for k, v in database.to_json().items() if k in DATABASE_KEYS - } - databases[database.id]["backend"] = database.backend - queries: dict[str, Any] = {} - - # These are unnecessary if sqllab backend persistence is disabled - if is_feature_enabled("SQLLAB_BACKEND_PERSISTENCE"): - # return all user queries associated with existing SQL editors - user_queries = ( - db.session.query(Query) - .filter_by(user_id=user_id) - .filter(Query.sql_editor_id.in_(tab_state_ids)) - .all() - ) - queries = { - query.client_id: dict(query.to_dict().items()) for query in user_queries - } - - return { - "tab_state_ids": tabs_state, - "active_tab": active_tab.to_dict() if active_tab else None, - "databases": databases, - "queries": queries, - } - @has_access @event_logger.log_this @expose( @@ -1076,9 +1016,8 @@ def _get_sqllab_tabs(user_id: int | None) -> dict[str, Any]: def sqllab(self) -> FlaskResponse: """SQL Editor""" payload = { - "defaultDbId": config["SQLLAB_DEFAULT_DBID"], "common": common_bootstrap_payload(g.user), - **self._get_sqllab_tabs(get_user_id()), + **bootstrap_sqllab_data(get_user_id()), } if form_data := request.form.get("form_data"): diff --git a/tests/integration_tests/core_tests.py b/tests/integration_tests/core_tests.py index e036602d0f..ddcfa189b2 100644 --- a/tests/integration_tests/core_tests.py +++ b/tests/integration_tests/core_tests.py @@ -52,6 +52,7 @@ from superset.models.slice import Slice from superset.models.sql_lab import Query from superset.result_set import SupersetResultSet +from superset.sqllab.utils import bootstrap_sqllab_data from superset.utils import core as utils from superset.utils.core import backend from superset.utils.database import get_example_database @@ -1135,7 +1136,8 @@ def test_sqllab_backend_persistence_payload(self): # we should have only 1 query returned, since the second one is not # associated with any tabs - payload = views.Superset._get_sqllab_tabs(user_id=user_id) + # TODO: replaces this spec by api/v1/sqllab spec later + payload = bootstrap_sqllab_data(user_id) self.assertEqual(len(payload["queries"]), 1) @mock.patch.dict(