Skip to content

Commit

Permalink
force limit only when there is no existing limit
Browse files Browse the repository at this point in the history
  • Loading branch information
timifasubaa committed May 25, 2018
1 parent c18ef89 commit 1aced9b
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 4 deletions.
4 changes: 3 additions & 1 deletion superset/sql_lab.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,7 @@ def handle_error(msg):
# Limit enforced only for retrieving the data, not for the CTA queries.
superset_query = SupersetQuery(rendered_query)
executed_sql = superset_query.stripped()
SQL_MAX_ROWS = int(app.config.get('SQL_MAX_ROW', None))
if not superset_query.is_select() and not database.allow_dml:
return handle_error(
'Only `SELECT` statements are allowed against this database')
Expand All @@ -185,7 +186,8 @@ def handle_error(msg):
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
elif (query.limit and superset_query.is_select()):
elif (not query.limit and superset_query.is_select() and SQL_MAX_ROWS):
query.limit = SQL_MAX_ROWS
executed_sql = database.apply_limit_to_sql(executed_sql, query.limit)
query.limit_used = True

Expand Down
14 changes: 14 additions & 0 deletions superset/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -882,3 +882,17 @@ def split_adhoc_filters_into_base_filters(fd):
fd['having_filters'] = simple_having_filters
fd['filters'] = simple_where_filters
del fd['adhoc_filters']


def get_limit_from_sql(sql):
sql = sql.lower()
limit = None
tokens = sql.split()
try:
if 'limit' in tokens:
limit_pos = tokens.index('limit') + 1
limit = int(tokens[limit_pos])
except Exception as e:
# fail quietly so we can get the more intelligible error from the database.
logging.error('Non-numeric limit added.\n{}'.format(e))
return limit
2 changes: 1 addition & 1 deletion superset/views/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -2394,7 +2394,7 @@ def sql_json(self):

query = Query(
database_id=int(database_id),
limit=int(app.config.get('SQL_MAX_ROW', None)),
limit=utils.get_limit_from_sql(sql),
sql=sql,
schema=schema,
select_as_cta=request.form.get('select_as_cta') == 'true',
Expand Down
2 changes: 0 additions & 2 deletions tests/celery_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -196,13 +196,11 @@ def test_run_async_query(self):
self.assertEqual([{'name': 'Admin'}], df.to_dict(orient='records'))
self.assertEqual(QueryStatus.SUCCESS, query.status)
self.assertTrue('FROM tmp_async_1' in query.select_sql)
self.assertTrue('LIMIT 666' in query.select_sql)
self.assertEqual(
'CREATE TABLE tmp_async_1 AS \nSELECT name FROM ab_role '
"WHERE name='Admin'", query.executed_sql)
self.assertEqual(sql_where, query.sql)
self.assertEqual(0, query.rows)
self.assertEqual(666, query.limit)
self.assertEqual(False, query.limit_used)
self.assertEqual(True, query.select_as_cta)
self.assertEqual(True, query.select_as_cta_used)
Expand Down

0 comments on commit 1aced9b

Please sign in to comment.