diff --git a/superset/assets/src/SqlLab/actions/sqlLab.js b/superset/assets/src/SqlLab/actions/sqlLab.js index 8f36ecad2..e9785f850 100644 --- a/superset/assets/src/SqlLab/actions/sqlLab.js +++ b/superset/assets/src/SqlLab/actions/sqlLab.js @@ -224,12 +224,12 @@ export function requestQueryResults(query) { return { type: REQUEST_QUERY_RESULTS, query }; } -export function fetchQueryResults(query) { +export function fetchQueryResults(query, displayLimit) { return function (dispatch) { dispatch(requestQueryResults(query)); return SupersetClient.get({ - endpoint: `/superset/results/${query.resultsKey}/`, + endpoint: `/superset/results/${query.resultsKey}/?rows=${displayLimit}`, parseMethod: 'text', }) .then(({ text = '{}' }) => { diff --git a/superset/assets/src/SqlLab/components/App.jsx b/superset/assets/src/SqlLab/components/App.jsx index c4f2960cc..0124f8d8c 100644 --- a/superset/assets/src/SqlLab/components/App.jsx +++ b/superset/assets/src/SqlLab/components/App.jsx @@ -100,7 +100,13 @@ class App extends React.PureComponent { render() { let content; if (this.state.hash) { - content = ; + content = ( + + ); } else { content = ( @@ -121,12 +127,14 @@ class App extends React.PureComponent { App.propTypes = { actions: PropTypes.object, localStorageUsageInKilobytes: PropTypes.number.isRequired, + common: PropTypes.object, }; function mapStateToProps(state) { - const { localStorageUsageInKilobytes } = state; + const { localStorageUsageInKilobytes, common } = state; return { localStorageUsageInKilobytes, + common, }; } diff --git a/superset/assets/src/SqlLab/components/QueryHistory.jsx b/superset/assets/src/SqlLab/components/QueryHistory.jsx index 41086237e..738ee901b 100644 --- a/superset/assets/src/SqlLab/components/QueryHistory.jsx +++ b/superset/assets/src/SqlLab/components/QueryHistory.jsx @@ -26,6 +26,7 @@ import QueryTable from './QueryTable'; const propTypes = { queries: PropTypes.array.isRequired, actions: PropTypes.object.isRequired, + displayLimit: PropTypes.number.isRequired, }; const QueryHistory = (props) => { @@ -38,6 +39,7 @@ const QueryHistory = (props) => { ]} queries={props.queries} actions={props.actions} + displayLimit={props.displayLimit} /> ); } diff --git a/superset/assets/src/SqlLab/components/QuerySearch.jsx b/superset/assets/src/SqlLab/components/QuerySearch.jsx index fcf2f6f71..97d2ae53f 100644 --- a/superset/assets/src/SqlLab/components/QuerySearch.jsx +++ b/superset/assets/src/SqlLab/components/QuerySearch.jsx @@ -37,6 +37,7 @@ import AsyncSelect from '../../components/AsyncSelect'; const propTypes = { actions: PropTypes.object.isRequired, height: PropTypes.string.isRequired, + displayLimit: PropTypes.number.isRequired, }; class QuerySearch extends React.PureComponent { @@ -270,6 +271,7 @@ class QuerySearch extends React.PureComponent { onDbClicked={this.onDbClicked} queries={this.state.queriesArray} actions={this.props.actions} + displayLimit={this.props.displayLimit} /> )} diff --git a/superset/assets/src/SqlLab/components/QueryTable.jsx b/superset/assets/src/SqlLab/components/QueryTable.jsx index 5cf122d03..983d8cfcc 100644 --- a/superset/assets/src/SqlLab/components/QueryTable.jsx +++ b/superset/assets/src/SqlLab/components/QueryTable.jsx @@ -37,6 +37,7 @@ const propTypes = { queries: PropTypes.array, onUserClicked: PropTypes.func, onDbClicked: PropTypes.func, + displayLimit: PropTypes.number.isRequired, }; const defaultProps = { columns: ['started', 'duration', 'rows'], @@ -81,8 +82,8 @@ class QueryTable extends React.PureComponent { openQueryInNewTab(query) { this.props.actions.cloneQueryToNewTab(query); } - openAsyncResults(query) { - this.props.actions.fetchQueryResults(query); + openAsyncResults(query, displayLimit) { + this.props.actions.fetchQueryResults(query, displayLimit); } clearQueryResults(query) { this.props.actions.clearQueryResults(query); @@ -151,10 +152,16 @@ class QueryTable extends React.PureComponent { } modalTitle={t('Data preview')} - beforeOpen={this.openAsyncResults.bind(this, query)} + beforeOpen={this.openAsyncResults.bind(this, query, this.props.displayLimit)} onExit={this.clearQueryResults.bind(this, query)} modalBody={ - + } /> ); diff --git a/superset/assets/src/SqlLab/components/ResultSet.jsx b/superset/assets/src/SqlLab/components/ResultSet.jsx index 1ebe812c1..b0f5c7e42 100644 --- a/superset/assets/src/SqlLab/components/ResultSet.jsx +++ b/superset/assets/src/SqlLab/components/ResultSet.jsx @@ -40,6 +40,7 @@ const propTypes = { cache: PropTypes.bool, height: PropTypes.number.isRequired, database: PropTypes.object, + displayLimit: PropTypes.number.isRequired, }; const defaultProps = { search: true, @@ -105,7 +106,7 @@ export default class ResultSet extends React.PureComponent { this.setState({ searchText: event.target.value }); } fetchResults(query) { - this.props.actions.fetchQueryResults(query); + this.props.actions.fetchQueryResults(query, this.props.displayLimit); } reFetchQueryResults(query) { this.props.actions.reFetchQueryResults(query); diff --git a/superset/assets/src/SqlLab/components/SouthPane.jsx b/superset/assets/src/SqlLab/components/SouthPane.jsx index debc8ded5..e68d05c9a 100644 --- a/superset/assets/src/SqlLab/components/SouthPane.jsx +++ b/superset/assets/src/SqlLab/components/SouthPane.jsx @@ -44,6 +44,7 @@ const propTypes = { height: PropTypes.number, databases: PropTypes.object.isRequired, offline: PropTypes.bool, + displayLimit: PropTypes.number.isRequired, }; const defaultProps = { @@ -97,6 +98,7 @@ export class SouthPane extends React.PureComponent { actions={props.actions} height={innerTabContentHeight} database={this.props.databases[latestQuery.dbId]} + displayLimit={this.props.displayLimit} /> ); } else { @@ -115,6 +117,7 @@ export class SouthPane extends React.PureComponent { actions={props.actions} cache height={innerTabContentHeight} + displayLimit={this.props.displayLimit} /> )); @@ -139,7 +142,11 @@ export class SouthPane extends React.PureComponent { title={t('Query History')} eventKey="History" > - + {dataPreviewTabs} diff --git a/superset/assets/src/SqlLab/components/SqlEditor.jsx b/superset/assets/src/SqlLab/components/SqlEditor.jsx index 001d91da7..d4cc424a4 100644 --- a/superset/assets/src/SqlLab/components/SqlEditor.jsx +++ b/superset/assets/src/SqlLab/components/SqlEditor.jsx @@ -71,6 +71,7 @@ const propTypes = { hideLeftBar: PropTypes.bool, defaultQueryLimit: PropTypes.number.isRequired, maxRow: PropTypes.number.isRequired, + displayLimit: PropTypes.number.isRequired, saveQueryWarning: PropTypes.string, scheduleQueryWarning: PropTypes.string, }; @@ -327,6 +328,7 @@ class SqlEditor extends React.PureComponent { dataPreviewQueries={this.props.dataPreviewQueries} actions={this.props.actions} height={southPaneHeight} + displayLimit={this.props.displayLimit} /> ); diff --git a/superset/assets/src/SqlLab/components/TabbedSqlEditors.jsx b/superset/assets/src/SqlLab/components/TabbedSqlEditors.jsx index c50b8b11d..7b7f1773d 100644 --- a/superset/assets/src/SqlLab/components/TabbedSqlEditors.jsx +++ b/superset/assets/src/SqlLab/components/TabbedSqlEditors.jsx @@ -32,6 +32,7 @@ import TabStatusIcon from './TabStatusIcon'; const propTypes = { actions: PropTypes.object.isRequired, defaultDbId: PropTypes.number, + displayLimit: PropTypes.number, defaultQueryLimit: PropTypes.number.isRequired, maxRow: PropTypes.number.isRequired, databases: PropTypes.object.isRequired, @@ -251,6 +252,7 @@ class TabbedSqlEditors extends React.PureComponent { hideLeftBar={this.state.hideLeftBar} defaultQueryLimit={this.props.defaultQueryLimit} maxRow={this.props.maxRow} + displayLimit={this.props.displayLimit} saveQueryWarning={this.props.saveQueryWarning} scheduleQueryWarning={this.props.scheduleQueryWarning} /> @@ -293,6 +295,7 @@ function mapStateToProps({ sqlLab, common }) { tabHistory: sqlLab.tabHistory, tables: sqlLab.tables, defaultDbId: sqlLab.defaultDbId, + displayLimit: common.conf.DISPLAY_MAX_ROW, offline: sqlLab.offline, defaultQueryLimit: common.conf.DEFAULT_SQLLAB_LIMIT, maxRow: common.conf.SQL_MAX_ROW, diff --git a/superset/views/base.py b/superset/views/base.py index fb6eccedc..07346d14f 100644 --- a/superset/views/base.py +++ b/superset/views/base.py @@ -47,6 +47,7 @@ FRONTEND_CONF_KEYS = ( "SQL_MAX_ROW", "SUPERSET_WEBSERVER_DOMAINS", "SQLLAB_SAVE_WARNING_MESSAGE", + "DISPLAY_MAX_ROW", ) diff --git a/superset/views/core.py b/superset/views/core.py index 0c6b6ec4c..69a4ff1e5 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2408,7 +2408,11 @@ class Superset(BaseSupersetView): @expose("/results//") @event_logger.log_this def results(self, key): - """Serves a key off of the results backend""" + """Serves a key off of the results backend + + It is possible to pass the `rows` query argument to limit the number + of rows returned. + """ if not results_backend: return json_error_response("Results backend isn't configured") @@ -2442,12 +2446,15 @@ class Superset(BaseSupersetView): payload = utils.zlib_decompress(blob, decode=not results_backend_use_msgpack) obj = _deserialize_results_payload(payload, query, results_backend_use_msgpack) + if "rows" in request.args: + try: + rows = int(request.args["rows"]) + except ValueError: + return json_error_response("Invalid `rows` argument", status=400) + obj = apply_display_max_row_limit(obj, rows) + return json_success( - json.dumps( - apply_display_max_row_limit(obj), - default=utils.json_iso_dttm_ser, - ignore_nan=True, - ) + json.dumps(obj, default=utils.json_iso_dttm_ser, ignore_nan=True) ) @has_access_api diff --git a/superset/views/utils.py b/superset/views/utils.py index a785600d0..231373537 100644 --- a/superset/views/utils.py +++ b/superset/views/utils.py @@ -175,7 +175,9 @@ def get_datasource_info( return datasource_id, datasource_type -def apply_display_max_row_limit(sql_results: Dict[str, Any]) -> Dict[str, Any]: +def apply_display_max_row_limit( + sql_results: Dict[str, Any], rows: Optional[int] = None +) -> Dict[str, Any]: """ Given a `sql_results` nested structure, applies a limit to the number of rows @@ -187,7 +189,7 @@ def apply_display_max_row_limit(sql_results: Dict[str, Any]) -> Dict[str, Any]: :param sql_results: The results of a sql query from sql_lab.get_sql_results :returns: The mutated sql_results structure """ - display_limit = app.config.get("DISPLAY_MAX_ROW") + display_limit = rows or app.config.get("DISPLAY_MAX_ROW") if ( display_limit and sql_results["status"] == QueryStatus.SUCCESS diff --git a/tests/core_tests.py b/tests/core_tests.py index 2110cdb0a..e2cf010b6 100644 --- a/tests/core_tests.py +++ b/tests/core_tests.py @@ -32,7 +32,7 @@ import pandas as pd import psycopg2 import sqlalchemy as sqla -from superset import dataframe, db, jinja_context, security_manager, sql_lab +from superset import app, dataframe, db, jinja_context, security_manager, sql_lab from superset.connectors.sqla.models import SqlaTable from superset.db_engine_specs.base import BaseEngineSpec from superset.db_engine_specs.mssql import MssqlEngineSpec @@ -756,6 +756,49 @@ class CoreTests(SupersetTestCase): resp = self.get_resp(f"/superset/select_star/{examples_db.id}/birth_names") self.assertIn("gender", resp) + @mock.patch("superset.views.core.results_backend_use_msgpack", False) + @mock.patch("superset.views.core.results_backend") + @mock.patch("superset.views.core.db") + def test_display_limit(self, mock_superset_db, mock_results_backend): + query_mock = mock.Mock() + query_mock.sql = "SELECT *" + query_mock.database = 1 + query_mock.schema = "superset" + mock_superset_db.session.query().filter_by().one_or_none.return_value = ( + query_mock + ) + + data = [{"col_0": i} for i in range(100)] + payload = { + "status": utils.QueryStatus.SUCCESS, + "query": {"rows": 100}, + "data": data, + } + # do not apply msgpack serialization + use_msgpack = app.config["RESULTS_BACKEND_USE_MSGPACK"] + app.config["RESULTS_BACKEND_USE_MSGPACK"] = False + serialized_payload = sql_lab._serialize_payload(payload, False) + compressed = utils.zlib_compress(serialized_payload) + mock_results_backend.get.return_value = compressed + + # get all results + result = json.loads(self.get_resp("/superset/results/key/")) + expected = {"status": "success", "query": {"rows": 100}, "data": data} + self.assertEqual(result, expected) + + # limit results to 1 + limited_data = data[:1] + result = json.loads(self.get_resp("/superset/results/key/?rows=1")) + expected = { + "status": "success", + "query": {"rows": 100}, + "data": limited_data, + "displayLimitReached": True, + } + self.assertEqual(result, expected) + + app.config["RESULTS_BACKEND_USE_MSGPACK"] = use_msgpack + def test_results_default_deserialization(self): use_new_deserialization = False data = [("a", 4, 4.0, "2019-08-18T16:39:16.660000")]