Skip to content

Commit

Permalink
refactor move apply_display_max_row_limit to sqllab package
Browse files Browse the repository at this point in the history
  • Loading branch information
ofekisr committed Sep 26, 2021
1 parent c596d42 commit 9b00d1a
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 4 deletions.
4 changes: 2 additions & 2 deletions superset/sqllab/command.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@
from superset.utils import core as utils
from superset.utils.dates import now_as_float
from superset.utils.sqllab_execution_context import SqlJsonExecutionContext
from superset.views.utils import apply_display_max_row_limit
from superset.sqllab.utils import apply_display_max_row_configuration_if_require

config = app.config
logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -397,7 +397,7 @@ def _to_payload_results_based( # pylint: disable=no-self-use
) -> str:
display_max_row = config["DISPLAY_MAX_ROW"]
return json.dumps(
apply_display_max_row_limit(execution_result, display_max_row),
apply_display_max_row_configuration_if_require(execution_result, display_max_row),
default=utils.pessimistic_json_iso_dttm_ser,
ignore_nan=True,
encoding=None,
Expand Down
47 changes: 47 additions & 0 deletions superset/sqllab/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from typing import Any, Dict

from superset.common.db_query_status import QueryStatus


def apply_display_max_row_configuration_if_require(
sql_results: Dict[str, Any], max_rows_in_result: int
) -> Dict[str, Any]:
"""
Given a `sql_results` nested structure, applies a limit to the number of rows
`sql_results` here is the nested structure coming out of sql_lab.get_sql_results, it
contains metadata about the query, as well as the data set returned by the query.
This method limits the number of rows adds a `displayLimitReached: True` flag to the
metadata.
:param max_rows_in_result:
:param sql_results: The results of a sql query from sql_lab.get_sql_results
:returns: The mutated sql_results structure
"""

def is_require_to_apply() -> bool:
return (
sql_results["status"] == QueryStatus.SUCCESS
and sql_results["query"]["rows"] > max_rows_in_result
)

if is_require_to_apply():
sql_results["data"] = sql_results["data"][:max_rows_in_result]
sql_results["displayLimitReached"] = True
return sql_results
4 changes: 2 additions & 2 deletions superset/views/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@
from superset.security.analytics_db_safety import check_sqlalchemy_uri
from superset.sql_parse import ParsedQuery, Table
from superset.sql_validators import get_validator_by_name
from superset.sqllab.utils import apply_display_max_row_configuration_if_require
from superset.sqllab.command import CommandResult, ExecuteSqlCommand
from superset.sqllab.command_status import SqlJsonExecutionStatus
from superset.tasks.async_queries import load_explore_json_into_cache
Expand Down Expand Up @@ -128,7 +129,6 @@
)
from superset.views.utils import (
_deserialize_results_payload,
apply_display_max_row_limit,
bootstrap_user_data,
check_datasource_perms,
check_explore_cache_perms,
Expand Down Expand Up @@ -2314,7 +2314,7 @@ def results_exec(key: str) -> FlaskResponse:
status=400,
) from ex

obj = apply_display_max_row_limit(obj, rows)
obj = apply_display_max_row_configuration_if_require(obj, rows)

return json_success(
json.dumps(
Expand Down

0 comments on commit 9b00d1a

Please sign in to comment.