Skip to content

Commit

Permalink
refactor sql_json view endpoint: separate set the query's limiting co…
Browse files Browse the repository at this point in the history
…ncern into ad hod method (apache#16653)
  • Loading branch information
ofekisr authored and Emmanuel Bavoux committed Nov 14, 2021
1 parent ee8aed6 commit df21b79
Showing 1 changed file with 36 additions and 14 deletions.
50 changes: 36 additions & 14 deletions superset/views/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -2600,20 +2600,7 @@ def _run_sql_json_exec_from_scratch(
self._validate_access(query, session)
rendered_query = self._render_query(execution_context, query, session)

if not (config.get("SQLLAB_CTAS_NO_LIMIT") and execution_context.select_as_cta):
# set LIMIT after template processing
db_engine_spec = execution_context.database.db_engine_spec # type: ignore
limits = [
db_engine_spec.get_limit_from_sql(rendered_query),
execution_context.limit,
]
if limits[0] is None or limits[0] > limits[1]: # type: ignore
query.limiting_factor = LimitingFactor.DROPDOWN
elif limits[1] > limits[0]: # type: ignore
query.limiting_factor = LimitingFactor.QUERY
else: # limits[0] == limits[1]
query.limiting_factor = LimitingFactor.QUERY_AND_DROPDOWN
query.limit = min(lim for lim in limits if lim is not None)
self._set_query_limit_if_required(execution_context, query, rendered_query)

# Flag for whether or not to expand data
# (feature that will expand Presto row objects and arrays)
Expand All @@ -2629,6 +2616,41 @@ def _run_sql_json_exec_from_scratch(
session, rendered_query, query, expand_data, log_params
)

def _set_query_limit_if_required(
self,
execution_context: SqlJsonExecutionContext,
query: Query,
rendered_query: str,
) -> None:
if self._is_required_to_set_limit(execution_context):
self._set_query_limit(execution_context, query, rendered_query)

def _set_query_limit( # pylint: disable=no-self-use
self,
execution_context: SqlJsonExecutionContext,
query: Query,
rendered_query: str,
) -> None:
db_engine_spec = execution_context.database.db_engine_spec # type: ignore
limits = [
db_engine_spec.get_limit_from_sql(rendered_query),
execution_context.limit,
]
if limits[0] is None or limits[0] > limits[1]: # type: ignore
query.limiting_factor = LimitingFactor.DROPDOWN
elif limits[1] > limits[0]: # type: ignore
query.limiting_factor = LimitingFactor.QUERY
else: # limits[0] == limits[1]
query.limiting_factor = LimitingFactor.QUERY_AND_DROPDOWN
query.limit = min(lim for lim in limits if lim is not None)

def _is_required_to_set_limit( # pylint: disable=no-self-use
self, execution_context: SqlJsonExecutionContext
) -> bool:
return not (
config.get("SQLLAB_CTAS_NO_LIMIT") and execution_context.select_as_cta
)

def _render_query( # pylint: disable=no-self-use
self, execution_context: SqlJsonExecutionContext, query: Query, session: Session
) -> str:
Expand Down

0 comments on commit df21b79

Please sign in to comment.