Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: preventing sql lab None limit value #17155

Merged
merged 4 commits into from
Oct 20, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 14 additions & 6 deletions superset/sql_lab.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetErrorException, SupersetErrorsException
from superset.extensions import celery_app
from superset.models.core import Database
from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
from superset.sql_parse import CtasMethod, ParsedQuery
Expand Down Expand Up @@ -221,12 +222,7 @@ def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-locals
):
if SQL_MAX_ROW and (not query.limit or query.limit > SQL_MAX_ROW):
query.limit = SQL_MAX_ROW
if query.limit:
# We are fetching one more than the requested limit in order
# to test whether there are more rows than the limit.
# Later, the extra row will be dropped before sending
# the results back to the user.
sql = database.apply_limit_to_sql(sql, increased_limit, force=True)
sql = apply_limit_if_exists(database, increased_limit, query, sql)

# Hook to allow environment-specific mutation (usually comments) to the SQL
sql = SQL_QUERY_MUTATOR(sql, user_name, security_manager, database)
Expand Down Expand Up @@ -291,6 +287,18 @@ def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-locals
return SupersetResultSet(data, cursor_description, db_engine_spec)


def apply_limit_if_exists(
database: Database, increased_limit: Optional[int], query: Query, sql: str
) -> str:
if query.limit and increased_limit:
# We are fetching one more than the requested limit in order
# to test whether there are more rows than the limit.
# Later, the extra row will be dropped before sending
# the results back to the user.
sql = database.apply_limit_to_sql(sql, increased_limit, force=True)
return sql


def _serialize_payload(
payload: Dict[Any, Any], use_msgpack: Optional[bool] = False
) -> Union[bytes, str]:
Expand Down
24 changes: 24 additions & 0 deletions tests/integration_tests/sqllab_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
execute_sql_statement,
get_sql_results,
SqlLabException,
apply_limit_if_exists,
)
from superset.sql_parse import CtasMethod
from superset.utils.core import (
Expand Down Expand Up @@ -990,6 +991,29 @@ def test_sql_json_soft_timeout(self):
]
}

def test_apply_limit_if_exists_when_incremented_limit_is_none(self):
sql = """
SET @value = 42;
SELECT @value AS foo;
"""
database = get_example_database()
mock_query = mock.MagicMock()
mock_query.limit = 300
final_sql = apply_limit_if_exists(database, None, mock_query, sql)

assert final_sql == sql

def test_apply_limit_if_exists_when_increased_limit(self):
sql = """
SET @value = 42;
SELECT @value AS foo;
"""
database = get_example_database()
mock_query = mock.MagicMock()
mock_query.limit = 300
final_sql = apply_limit_if_exists(database, 1000, mock_query, sql)
assert "LIMIT 1000" in final_sql


@pytest.mark.parametrize("spec", [HiveEngineSpec, PrestoEngineSpec])
def test_cancel_query_implicit(spec: BaseEngineSpec) -> None:
Expand Down