superset-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From grace...@apache.org
Subject [incubator-superset] branch master updated: Add separate limit setting for SqlLab (#4941)
Date Wed, 07 Nov 2018 23:57:52 GMT
This is an automated email from the ASF dual-hosted git repository.

graceguo pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 0584e36  Add separate limit setting for SqlLab (#4941)
0584e36 is described below

commit 0584e3629feaad17dc1391760aeb8a0cf6e8832f
Author: Jeffrey Wang <jeffreywang93@gmail.com>
AuthorDate: Wed Nov 7 18:57:44 2018 -0500

    Add separate limit setting for SqlLab (#4941)
    
    * Add separate limit setting for SqlLab
    
    Use separate param for wrap sql
    
    Get query limit from config
    
    unit tests for limit control rendering in sql editor
    
    py unit test
    
    pg tests
    
    Add max rows limit
    
    Remove concept of infinity, always require defined limits
    
    consistency
    
    Assert on validation errors instead of tooltip
    
    fix unit tests
    
    attempt persist state
    
    pr comments and linting
    
    * load configs in via common param
    
    * default to 1k
---
 .../spec/javascripts/sqllab/LimitControl_spec.jsx  |  63 +++++++++++
 .../spec/javascripts/sqllab/SqlEditor_spec.jsx     |  19 +++-
 .../javascripts/sqllab/TabbedSqlEditors_spec.jsx   |   2 +
 .../assets/spec/javascripts/sqllab/fixtures.js     |   6 +
 .../javascripts/sqllab/reducers/sqlLab_spec.js     |   5 +
 superset/assets/src/SqlLab/actions/sqlLab.js       |   7 ++
 .../assets/src/SqlLab/components/LimitControl.jsx  | 126 +++++++++++++++++++++
 .../assets/src/SqlLab/components/SqlEditor.jsx     |  20 +++-
 .../src/SqlLab/components/TabbedSqlEditors.jsx     |   8 +-
 .../assets/src/SqlLab/reducers/getInitialState.js  |   1 +
 superset/assets/src/SqlLab/reducers/sqlLab.js      |   5 +
 superset/config.py                                 |   4 +-
 superset/sql_lab.py                                |   9 +-
 superset/views/base.py                             |   2 +
 superset/views/core.py                             |  12 +-
 tests/base_tests.py                                |   5 +-
 tests/sqllab_tests.py                              |  23 ++++
 17 files changed, 303 insertions(+), 14 deletions(-)

diff --git a/superset/assets/spec/javascripts/sqllab/LimitControl_spec.jsx b/superset/assets/spec/javascripts/sqllab/LimitControl_spec.jsx
new file mode 100644
index 0000000..7869eae
--- /dev/null
+++ b/superset/assets/spec/javascripts/sqllab/LimitControl_spec.jsx
@@ -0,0 +1,63 @@
+import React from 'react';
+import { shallow } from 'enzyme';
+
+import { Label } from 'react-bootstrap';
+import LimitControl from '../../../src/SqlLab/components/LimitControl';
+import ControlHeader from '../../../src/explore/components/ControlHeader';
+
+describe('LimitControl', () => {
+  const defaultProps = {
+    value: 0,
+    defaultQueryLimit: 1000,
+    maxRow: 100000,
+    onChange: () => {},
+  };
+  let wrapper;
+  const factory = o => <LimitControl {...o} />;
+  beforeEach(() => {
+    wrapper = shallow(factory(defaultProps));
+  });
+  it('is a valid element', () => {
+    expect(React.isValidElement(<LimitControl {...defaultProps} />)).toEqual(true);
+  });
+  it('renders a Label', () => {
+    expect(wrapper.find(Label)).toHaveLength(1);
+  });
+  it('loads the correct state', () => {
+    const value = 100;
+    wrapper = shallow(factory({ ...defaultProps, value }));
+    expect(wrapper.state().textValue).toEqual(value.toString());
+    wrapper.find(Label).first().simulate('click');
+    expect(wrapper.state().showOverlay).toBe(true);
+    expect(wrapper.find(ControlHeader).props().validationErrors).toHaveLength(0);
+  });
+  it('handles invalid value', () => {
+    wrapper.find(Label).first().simulate('click');
+    wrapper.setState({ textValue: 'invalid' });
+    expect(wrapper.find(ControlHeader).props().validationErrors).toHaveLength(1);
+  });
+  it('handles negative value', () => {
+    wrapper.find(Label).first().simulate('click');
+    wrapper.setState({ textValue: '-1' });
+    expect(wrapper.find(ControlHeader).props().validationErrors).toHaveLength(1);
+  });
+  it('handles value above max row', () => {
+    wrapper.find(Label).first().simulate('click');
+    wrapper.setState({ textValue: (defaultProps.maxRow + 1).toString() });
+    expect(wrapper.find(ControlHeader).props().validationErrors).toHaveLength(1);
+  });
+  it('opens and closes', () => {
+    wrapper.find(Label).first().simulate('click');
+    expect(wrapper.state().showOverlay).toBe(true);
+    wrapper.find('.ok').first().simulate('click');
+    expect(wrapper.state().showOverlay).toBe(false);
+  });
+  it('resets and closes', () => {
+    const value = 100;
+    wrapper = shallow(factory({ ...defaultProps, value }));
+    wrapper.find(Label).first().simulate('click');
+    expect(wrapper.state().textValue).toEqual(value.toString());
+    wrapper.find('.reset').simulate('click');
+    expect(wrapper.state().textValue).toEqual(defaultProps.defaultQueryLimit.toString());
+  });
+});
diff --git a/superset/assets/spec/javascripts/sqllab/SqlEditor_spec.jsx b/superset/assets/spec/javascripts/sqllab/SqlEditor_spec.jsx
index 822192d..95db9b0 100644
--- a/superset/assets/spec/javascripts/sqllab/SqlEditor_spec.jsx
+++ b/superset/assets/spec/javascripts/sqllab/SqlEditor_spec.jsx
@@ -1,7 +1,8 @@
 import React from 'react';
 import { shallow } from 'enzyme';
 
-import { initialState, queries, table } from './fixtures';
+import { defaultQueryEditor, initialState, queries, table } from './fixtures';
+import LimitControl from '../../../src/SqlLab/components/LimitControl';
 import SqlEditor from '../../../src/SqlLab/components/SqlEditor';
 import SqlEditorLeftBar from '../../../src/SqlLab/components/SqlEditorLeftBar';
 
@@ -16,6 +17,8 @@ describe('SqlEditor', () => {
     getHeight: () => ('100px'),
     editorQueries: [],
     dataPreviewQueries: [],
+    defaultQueryLimit: 1000,
+    maxRow: 100000,
   };
   it('is valid', () => {
     expect(
@@ -26,4 +29,18 @@ describe('SqlEditor', () => {
     const wrapper = shallow(<SqlEditor {...mockedProps} />);
     expect(wrapper.find(SqlEditorLeftBar)).toHaveLength(1);
   });
+  it('render a LimitControl with default limit', () => {
+    const defaultQueryLimit = 101;
+    const updatedProps = { ...mockedProps, defaultQueryLimit };
+    const wrapper = shallow(<SqlEditor {...updatedProps} />);
+    expect(wrapper.find(LimitControl)).toHaveLength(1);
+    expect(wrapper.find(LimitControl).props().value).toEqual(defaultQueryLimit);
+  });
+  it('render a LimitControl with existing limit', () => {
+    const queryEditor = { ...defaultQueryEditor, queryLimit: 101 };
+    const updatedProps = { ...mockedProps, queryEditor };
+    const wrapper = shallow(<SqlEditor {...updatedProps} />);
+    expect(wrapper.find(LimitControl)).toHaveLength(1);
+    expect(wrapper.find(LimitControl).props().value).toEqual(queryEditor.queryLimit);
+  });
 });
diff --git a/superset/assets/spec/javascripts/sqllab/TabbedSqlEditors_spec.jsx b/superset/assets/spec/javascripts/sqllab/TabbedSqlEditors_spec.jsx
index 33d1e47..111aba9 100644
--- a/superset/assets/spec/javascripts/sqllab/TabbedSqlEditors_spec.jsx
+++ b/superset/assets/spec/javascripts/sqllab/TabbedSqlEditors_spec.jsx
@@ -52,6 +52,8 @@ describe('TabbedSqlEditors', () => {
     editorHeight: '',
     getHeight: () => ('100px'),
     database: {},
+    defaultQueryLimit: 1000,
+    maxRow: 100000,
   };
   const getWrapper = () => (
     shallow(<TabbedSqlEditors {...mockedProps} />, {
diff --git a/superset/assets/spec/javascripts/sqllab/fixtures.js b/superset/assets/spec/javascripts/sqllab/fixtures.js
index 02a6c12..99ba6a8 100644
--- a/superset/assets/spec/javascripts/sqllab/fixtures.js
+++ b/superset/assets/spec/javascripts/sqllab/fixtures.js
@@ -367,6 +367,12 @@ export const initialState = {
     activeSouthPaneTab: 'Results',
   },
   messageToasts: [],
+  common: {
+    conf: {
+      DEFAULT_SQLLAB_LIMIT: 1000,
+      SQL_MAX_ROW: 100000,
+    },
+  },
 };
 
 export const query = {
diff --git a/superset/assets/spec/javascripts/sqllab/reducers/sqlLab_spec.js b/superset/assets/spec/javascripts/sqllab/reducers/sqlLab_spec.js
index 964daa8..ba189ca 100644
--- a/superset/assets/spec/javascripts/sqllab/reducers/sqlLab_spec.js
+++ b/superset/assets/spec/javascripts/sqllab/reducers/sqlLab_spec.js
@@ -81,6 +81,11 @@ describe('sqlLabReducer', () => {
       newState = sqlLabReducer(newState, actions.queryEditorSetSql(qe, sql));
       expect(newState.queryEditors[1].sql).toBe(sql);
     });
+    it('should not fail while setting queryLimit', () => {
+      const queryLimit = 101;
+      newState = sqlLabReducer(newState, actions.queryEditorSetQueryLimit(qe, queryLimit));
+      expect(newState.queryEditors[1].queryLimit).toEqual(queryLimit);
+    });
     it('should set selectedText', () => {
       const selectedText = 'TEST';
       expect(newState.queryEditors[0].selectedText).toBeNull();
diff --git a/superset/assets/src/SqlLab/actions/sqlLab.js b/superset/assets/src/SqlLab/actions/sqlLab.js
index 17e022b..084aaea 100644
--- a/superset/assets/src/SqlLab/actions/sqlLab.js
+++ b/superset/assets/src/SqlLab/actions/sqlLab.js
@@ -27,6 +27,7 @@ export const QUERY_EDITOR_SET_SCHEMA = 'QUERY_EDITOR_SET_SCHEMA';
 export const QUERY_EDITOR_SET_TITLE = 'QUERY_EDITOR_SET_TITLE';
 export const QUERY_EDITOR_SET_AUTORUN = 'QUERY_EDITOR_SET_AUTORUN';
 export const QUERY_EDITOR_SET_SQL = 'QUERY_EDITOR_SET_SQL';
+export const QUERY_EDITOR_SET_QUERY_LIMIT = 'QUERY_EDITOR_SET_QUERY_LIMIT';
 export const QUERY_EDITOR_SET_TEMPLATE_PARAMS = 'QUERY_EDITOR_SET_TEMPLATE_PARAMS';
 export const QUERY_EDITOR_SET_SELECTED_TEXT = 'QUERY_EDITOR_SET_SELECTED_TEXT';
 export const QUERY_EDITOR_PERSIST_HEIGHT = 'QUERY_EDITOR_PERSIST_HEIGHT';
@@ -141,6 +142,7 @@ export function runQuery(query) {
       tmp_table_name: query.tempTableName,
       select_as_cta: query.ctas,
       templateParams: query.templateParams,
+      queryLimit: query.queryLimit,
     };
 
     return SupersetClient.post({
@@ -230,6 +232,10 @@ export function queryEditorSetSql(queryEditor, sql) {
   return { type: QUERY_EDITOR_SET_SQL, queryEditor, sql };
 }
 
+export function queryEditorSetQueryLimit(queryEditor, queryLimit) {
+  return { type: QUERY_EDITOR_SET_QUERY_LIMIT, queryEditor, queryLimit };
+}
+
 export function queryEditorSetTemplateParams(queryEditor, templateParams) {
   return { type: QUERY_EDITOR_SET_TEMPLATE_PARAMS, queryEditor, templateParams };
 }
@@ -325,6 +331,7 @@ export function reFetchQueryResults(query) {
       tab: '',
       runAsync: false,
       ctas: false,
+      queryLimit: query.queryLimit,
     };
     dispatch(runQuery(newQuery));
     dispatch(changeDataPreviewId(query.id, newQuery));
diff --git a/superset/assets/src/SqlLab/components/LimitControl.jsx b/superset/assets/src/SqlLab/components/LimitControl.jsx
new file mode 100644
index 0000000..cdb1f57
--- /dev/null
+++ b/superset/assets/src/SqlLab/components/LimitControl.jsx
@@ -0,0 +1,126 @@
+import React from 'react';
+import PropTypes from 'prop-types';
+import {
+  Button,
+  Label,
+  FormGroup,
+  FormControl,
+  Overlay,
+  Popover,
+} from 'react-bootstrap';
+import { t } from '@superset-ui/translation';
+
+import ControlHeader from '../../explore/components/ControlHeader';
+
+const propTypes = {
+  value: PropTypes.number,
+  defaultQueryLimit: PropTypes.number.isRequired,
+  maxRow: PropTypes.number.isRequired,
+  onChange: PropTypes.func.isRequired,
+};
+
+export default class LimitControl extends React.PureComponent {
+  constructor(props) {
+    super(props);
+    const { value, defaultQueryLimit } = props;
+    this.state = {
+      textValue: value.toString() || defaultQueryLimit.toString(),
+      showOverlay: false,
+    };
+    this.handleHide = this.handleHide.bind(this);
+    this.handleToggle = this.handleToggle.bind(this);
+    this.submitAndClose = this.submitAndClose.bind(this);
+  }
+
+  setValueAndClose(val) {
+    this.setState({ textValue: val }, this.submitAndClose);
+  }
+
+  submitAndClose() {
+    const value = parseInt(this.state.textValue, 10) || this.props.defaultQueryLimit;
+    this.props.onChange(value);
+    this.setState({ showOverlay: false });
+  }
+
+  isValidLimit(limit) {
+    const value =  parseInt(limit, 10);
+    return !(Number.isNaN(value) || value <= 0 || (this.props.maxRow && value
> this.props.maxRow));
+  }
+
+  handleToggle() {
+    this.setState({ showOverlay: !this.state.showOverlay });
+  }
+
+  handleHide() {
+    this.setState({ showOverlay: false });
+  }
+
+  renderPopover() {
+    const textValue = this.state.textValue;
+    const isValid = this.isValidLimit(textValue);
+    const errorMsg = 'Row limit must be positive integer' +
+      (this.props.maxRow ? ` and not greater than ${this.props.maxRow}` : '');
+    return (
+      <Popover id="sqllab-limit-results">
+        <div style={{ width: '100px' }}>
+          <ControlHeader
+            label={t('Row limit')}
+            validationErrors={!isValid ? [t(errorMsg)] : []}
+          />
+          <FormGroup>
+            <FormControl
+              type="text"
+              value={textValue}
+              placeholder={t(`Max: ${this.props.maxRow}`)}
+              bsSize="small"
+              onChange={e => this.setState({ textValue: e.target.value })}
+            />
+          </FormGroup>
+          <div className="clearfix">
+            <Button
+              bsSize="small"
+              bsStyle="primary"
+              className="float-left ok"
+              disabled={!isValid}
+              onClick={this.submitAndClose}
+            >
+              Ok
+            </Button>
+            <Button
+              bsSize="small"
+              className="float-right reset"
+              onClick={this.setValueAndClose.bind(this, this.props.defaultQueryLimit.toString())}
+            >
+              Reset
+            </Button>
+          </div>
+        </div>
+      </Popover>
+    );
+  }
+
+  render() {
+    return (
+      <div>
+        <Label
+          style={{ cursor: 'pointer' }}
+          onClick={this.handleToggle}
+        >
+          LIMIT {this.props.value || this.props.maxRow}
+        </Label>
+        <Overlay
+          rootClose
+          show={this.state.showOverlay}
+          onHide={this.handleHide}
+          trigger="click"
+          placement="right"
+          target={this}
+        >
+          {this.renderPopover()}
+        </Overlay>
+      </div>
+    );
+  }
+}
+
+LimitControl.propTypes = propTypes;
diff --git a/superset/assets/src/SqlLab/components/SqlEditor.jsx b/superset/assets/src/SqlLab/components/SqlEditor.jsx
index e739efa..7a38770 100644
--- a/superset/assets/src/SqlLab/components/SqlEditor.jsx
+++ b/superset/assets/src/SqlLab/components/SqlEditor.jsx
@@ -17,6 +17,7 @@ import SplitPane from 'react-split-pane';
 import { t } from '@superset-ui/translation';
 
 import Button from '../../components/Button';
+import LimitControl from './LimitControl';
 import TemplateParamsEditor from './TemplateParamsEditor';
 import SouthPane from './SouthPane';
 import SaveQuery from './SaveQuery';
@@ -38,6 +39,8 @@ const propTypes = {
   dataPreviewQueries: PropTypes.array.isRequired,
   queryEditor: PropTypes.object.isRequired,
   hideLeftBar: PropTypes.bool,
+  defaultQueryLimit: PropTypes.number.isRequired,
+  maxRow: PropTypes.number.isRequired,
 };
 
 const defaultProps = {
@@ -130,6 +133,9 @@ class SqlEditor extends React.PureComponent {
   setQueryEditorSql(sql) {
     this.props.actions.queryEditorSetSql(this.props.queryEditor, sql);
   }
+  setQueryLimit(queryLimit) {
+    this.props.actions.queryEditorSetQueryLimit(this.props.queryEditor, queryLimit);
+  }
   runQuery() {
     this.startQuery(!(this.props.database || {}).allow_run_sync);
   }
@@ -143,6 +149,7 @@ class SqlEditor extends React.PureComponent {
       schema: qe.schema,
       tempTableName: ctas ? this.state.ctas : '',
       templateParams: qe.templateParams,
+      queryLimit: qe.queryLimit,
       runAsync,
       ctas,
     };
@@ -236,7 +243,18 @@ class SqlEditor extends React.PureComponent {
             <span className="m-r-5">
               <ShareSqlLabQuery queryEditor={qe} />
             </span>
-            {ctasControls}
+            <span className="m-r-5">
+              {ctasControls}
+            </span>
+            <span className="inlineBlock m-r-5">
+              <LimitControl
+                value={(this.props.queryEditor.queryLimit !== undefined) ?
+                  this.props.queryEditor.queryLimit : this.props.defaultQueryLimit}
+                defaultQueryLimit={this.props.defaultQueryLimit}
+                maxRow={this.props.maxRow}
+                onChange={this.setQueryLimit.bind(this)}
+              />
+            </span>
             <span className="m-l-5">
               <Hotkeys
                 header="Keyboard shortcuts"
diff --git a/superset/assets/src/SqlLab/components/TabbedSqlEditors.jsx b/superset/assets/src/SqlLab/components/TabbedSqlEditors.jsx
index 0fb5f21..7407dcb 100644
--- a/superset/assets/src/SqlLab/components/TabbedSqlEditors.jsx
+++ b/superset/assets/src/SqlLab/components/TabbedSqlEditors.jsx
@@ -14,6 +14,8 @@ import TabStatusIcon from './TabStatusIcon';
 const propTypes = {
   actions: PropTypes.object.isRequired,
   defaultDbId: PropTypes.number,
+  defaultQueryLimit: PropTypes.number.isRequired,
+  maxRow: PropTypes.number.isRequired,
   databases: PropTypes.object.isRequired,
   queries: PropTypes.object.isRequired,
   queryEditors: PropTypes.array,
@@ -212,6 +214,8 @@ class TabbedSqlEditors extends React.PureComponent {
                   database={database}
                   actions={this.props.actions}
                   hideLeftBar={this.state.hideLeftBar}
+                  defaultQueryLimit={this.props.defaultQueryLimit}
+                  maxRow={this.props.maxRow}
                 />
               )}
             </div>
@@ -243,7 +247,7 @@ class TabbedSqlEditors extends React.PureComponent {
 TabbedSqlEditors.propTypes = propTypes;
 TabbedSqlEditors.defaultProps = defaultProps;
 
-function mapStateToProps({ sqlLab }) {
+function mapStateToProps({ sqlLab, common }) {
   return {
     databases: sqlLab.databases,
     queryEditors: sqlLab.queryEditors,
@@ -252,6 +256,8 @@ function mapStateToProps({ sqlLab }) {
     tables: sqlLab.tables,
     defaultDbId: sqlLab.defaultDbId,
     offline: sqlLab.offline,
+    defaultQueryLimit: common.conf.DEFAULT_SQLLAB_LIMIT,
+    maxRow: common.conf.SQL_MAX_ROW,
   };
 }
 function mapDispatchToProps(dispatch) {
diff --git a/superset/assets/src/SqlLab/reducers/getInitialState.js b/superset/assets/src/SqlLab/reducers/getInitialState.js
index 14bd677..3351a22 100644
--- a/superset/assets/src/SqlLab/reducers/getInitialState.js
+++ b/superset/assets/src/SqlLab/reducers/getInitialState.js
@@ -11,6 +11,7 @@ export default function getInitialState({ defaultDbId, ...restBootstrapData
}) {
     latestQueryId: null,
     autorun: false,
     dbId: defaultDbId,
+    queryLimit: restBootstrapData.common.conf.DEFAULT_SQLLAB_LIMIT,
   };
 
   return {
diff --git a/superset/assets/src/SqlLab/reducers/sqlLab.js b/superset/assets/src/SqlLab/reducers/sqlLab.js
index b8a27a9..ff881c3 100644
--- a/superset/assets/src/SqlLab/reducers/sqlLab.js
+++ b/superset/assets/src/SqlLab/reducers/sqlLab.js
@@ -32,6 +32,8 @@ export default function sqlLabReducer(state = {}, action) {
         schema: action.query.schema ? action.query.schema : null,
         autorun: true,
         sql: action.query.sql,
+        queryLimit: action.query.queryLimit,
+        maxRow: action.query.maxRow,
       };
 
       return sqlLabReducer(state, actions.addQueryEditor(qe));
@@ -204,6 +206,9 @@ export default function sqlLabReducer(state = {}, action) {
     [actions.QUERY_EDITOR_SET_SQL]() {
       return alterInArr(state, 'queryEditors', action.queryEditor, { sql: action.sql });
     },
+    [actions.QUERY_EDITOR_SET_QUERY_LIMIT]() {
+      return alterInArr(state, 'queryEditors', action.queryEditor, { queryLimit: action.queryLimit
});
+    },
     [actions.QUERY_EDITOR_SET_TEMPLATE_PARAMS]() {
       return alterInArr(state, 'queryEditors', action.queryEditor, {
         templateParams: action.templateParams,
diff --git a/superset/config.py b/superset/config.py
index c943739..a34e5d9 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -282,8 +282,8 @@ MAPBOX_API_KEY = os.environ.get('MAPBOX_API_KEY', '')
 # in the results backend. This also becomes the limit when exporting CSVs
 SQL_MAX_ROW = 100000
 
-# Limit to be returned to the frontend.
-DISPLAY_MAX_ROW = 1000
+# Default row limit for SQL Lab queries
+DEFAULT_SQLLAB_LIMIT = 1000
 
 # Maximum number of tables/views displayed in the dropdown window in SQL Lab.
 MAX_TABLE_NAMES = 3000
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index d211f02..1d4171b 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -150,10 +150,11 @@ def execute_sql(
                 query.user_id, start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
         executed_sql = superset_query.as_create_table(query.tmp_table_name)
         query.select_as_cta_used = True
-    if (superset_query.is_select() and SQL_MAX_ROWS and
-            (not query.limit or query.limit > SQL_MAX_ROWS)):
-        query.limit = SQL_MAX_ROWS
-        executed_sql = database.apply_limit_to_sql(executed_sql, query.limit)
+    if superset_query.is_select():
+        if SQL_MAX_ROWS and (not query.limit or query.limit > SQL_MAX_ROWS):
+            query.limit = SQL_MAX_ROWS
+        if query.limit:
+            executed_sql = database.apply_limit_to_sql(executed_sql, query.limit)
 
     # Hook to allow environment-specific mutation (usually comments) to the SQL
     SQL_QUERY_MUTATOR = config.get('SQL_QUERY_MUTATOR')
diff --git a/superset/views/base.py b/superset/views/base.py
index 4cba8b2..4b4926c 100644
--- a/superset/views/base.py
+++ b/superset/views/base.py
@@ -24,6 +24,8 @@ FRONTEND_CONF_KEYS = (
     'SUPERSET_WEBSERVER_TIMEOUT',
     'SUPERSET_DASHBOARD_POSITION_DATA_LIMIT',
     'ENABLE_JAVASCRIPT_CONTROLS',
+    'DEFAULT_SQLLAB_LIMIT',
+    'SQL_MAX_ROW',
 )
 
 
diff --git a/superset/views/core.py b/superset/views/core.py
index a16bceb..1f11e70 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -2456,7 +2456,7 @@ class Superset(BaseSupersetView):
                 '{}'.format(rejected_tables)), status=403)
 
         payload = utils.zlib_decompress_to_string(blob)
-        display_limit = app.config.get('DISPLAY_MAX_ROW', None)
+        display_limit = app.config.get('DEFAULT_SQLLAB_LIMIT', None)
         if display_limit:
             payload_json = json.loads(payload)
             payload_json['data'] = payload_json['data'][:display_limit]
@@ -2495,6 +2495,12 @@ class Superset(BaseSupersetView):
         schema = request.form.get('schema') or None
         template_params = json.loads(
             request.form.get('templateParams') or '{}')
+        limit = int(request.form.get('queryLimit', 0))
+        if limit < 0:
+            logging.warning(
+                'Invalid limit of {} specified. Defaulting to max limit.'.format(limit))
+            limit = 0
+        limit = limit or app.config.get('SQL_MAX_ROW')
 
         session = db.session()
         mydb = session.query(models.Database).filter_by(id=database_id).first()
@@ -2520,10 +2526,10 @@ class Superset(BaseSupersetView):
             )
 
         client_id = request.form.get('client_id') or utils.shortid()[:10]
-
+        limits = [mydb.db_engine_spec.get_limit_from_sql(sql), limit]
         query = Query(
             database_id=int(database_id),
-            limit=mydb.db_engine_spec.get_limit_from_sql(sql),
+            limit=min(lim for lim in limits if lim is not None),
             sql=sql,
             schema=schema,
             select_as_cta=request.form.get('select_as_cta') == 'true',
diff --git a/tests/base_tests.py b/tests/base_tests.py
index b9b4649..7ff1036 100644
--- a/tests/base_tests.py
+++ b/tests/base_tests.py
@@ -154,7 +154,8 @@ class SupersetTestCase(unittest.TestCase):
                     perm.view_menu and table.perm in perm.view_menu.name):
                 security_manager.del_permission_role(public_role, perm)
 
-    def run_sql(self, sql, client_id=None, user_name=None, raise_on_error=False):
+    def run_sql(self, sql, client_id=None, user_name=None, raise_on_error=False,
+                query_limit=None):
         if user_name:
             self.logout()
             self.login(username=(user_name if user_name else 'admin'))
@@ -163,7 +164,7 @@ class SupersetTestCase(unittest.TestCase):
             '/superset/sql_json/',
             raise_on_error=False,
             data=dict(database_id=dbid, sql=sql, select_as_create_as=False,
-                      client_id=client_id),
+                      client_id=client_id, queryLimit=query_limit),
         )
         if raise_on_error and 'error' in resp:
             raise Exception('run_sql failed')
diff --git a/tests/sqllab_tests.py b/tests/sqllab_tests.py
index ca2063c..95e3fbc 100644
--- a/tests/sqllab_tests.py
+++ b/tests/sqllab_tests.py
@@ -260,6 +260,29 @@ class SqlLabTests(SupersetTestCase):
         resp = self.get_json_resp('/superset/sqllab_viz/', data=data)
         self.assertIn('table_id', resp)
 
+    def test_sql_limit(self):
+        self.login('admin')
+        test_limit = 1
+        data = self.run_sql(
+            'SELECT * FROM ab_user',
+            client_id='sql_limit_1')
+        self.assertGreater(len(data['data']), test_limit)
+        data = self.run_sql(
+            'SELECT * FROM ab_user',
+            client_id='sql_limit_2',
+            query_limit=test_limit)
+        self.assertEquals(len(data['data']), test_limit)
+        data = self.run_sql(
+            'SELECT * FROM ab_user LIMIT {}'.format(test_limit),
+            client_id='sql_limit_3',
+            query_limit=test_limit + 1)
+        self.assertEquals(len(data['data']), test_limit)
+        data = self.run_sql(
+            'SELECT * FROM ab_user LIMIT {}'.format(test_limit + 1),
+            client_id='sql_limit_4',
+            query_limit=test_limit)
+        self.assertEquals(len(data['data']), test_limit)
+
 
 if __name__ == '__main__':
     unittest.main()


Mime
View raw message