refactor sql_json view endpoint: separate concern into ad hod method (#16595)

This commit is contained in:
ofekisr
2021-09-05 23:18:17 +03:00
committed by GitHub
parent be77ad2288
commit e60b489867
3 changed files with 17 additions and 14 deletions

View File

@@ -105,7 +105,7 @@ class SqlJsonExecutionContext: # pylint: disable=too-many-instance-attributes
limit = 0 limit = 0
return limit return limit
def _get_user_id(self) -> Optional[int]: # pylint: disable=R0201 def _get_user_id(self) -> Optional[int]: # pylint: disable=no-self-use
try: try:
return g.user.get_id() if g.user else None return g.user.get_id() if g.user else None
except RuntimeError: except RuntimeError:
@@ -135,7 +135,7 @@ class SqlJsonExecutionContext: # pylint: disable=too-many-instance-attributes
pass pass
def create_query(self) -> Query: def create_query(self) -> Query:
# pylint: disable=C0301 # pylint: disable=line-too-long
start_time = now_as_float() start_time = now_as_float()
if self.select_as_cta: if self.select_as_cta:
return Query( return Query(
@@ -167,7 +167,7 @@ class SqlJsonExecutionContext: # pylint: disable=too-many-instance-attributes
) )
class CreateTableAsSelect: # pylint: disable=R0903 class CreateTableAsSelect: # pylint: disable=too-few-public-methods
ctas_method: CtasMethod ctas_method: CtasMethod
target_schema_name: Optional[str] target_schema_name: Optional[str]
target_table_name: str target_table_name: str

View File

@@ -2556,7 +2556,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
"user_agent": cast(Optional[str], request.headers.get("USER_AGENT")) "user_agent": cast(Optional[str], request.headers.get("USER_AGENT"))
} }
execution_context = SqlJsonExecutionContext(request.json) execution_context = SqlJsonExecutionContext(request.json)
return self.sql_json_exec(execution_context, request.json, log_params) return self.sql_json_exec(execution_context, log_params)
@classmethod @classmethod
def is_query_handled(cls, query: Optional[Query]) -> bool: def is_query_handled(cls, query: Optional[Query]) -> bool:
@@ -2569,7 +2569,6 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
def sql_json_exec( # pylint: disable=too-many-statements def sql_json_exec( # pylint: disable=too-many-statements
self, self,
execution_context: SqlJsonExecutionContext, execution_context: SqlJsonExecutionContext,
query_params: Dict[str, Any],
log_params: Optional[Dict[str, Any]] = None, log_params: Optional[Dict[str, Any]] = None,
) -> FlaskResponse: ) -> FlaskResponse:
"""Runs arbitrary sql and returns data as json""" """Runs arbitrary sql and returns data as json"""
@@ -2582,6 +2581,16 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
payload = self._convert_query_to_payload(cast(Query, query)) payload = self._convert_query_to_payload(cast(Query, query))
return json_success(payload) return json_success(payload)
return self._run_sql_json_exec_from_scratch(
execution_context, session, log_params
)
def _run_sql_json_exec_from_scratch(
self,
execution_context: SqlJsonExecutionContext,
session: Session,
log_params: Optional[Dict[str, Any]] = None,
) -> FlaskResponse:
execution_context.set_database( execution_context.set_database(
self._get_the_query_db(execution_context, session) self._get_the_query_db(execution_context, session)
) )
@@ -2656,10 +2665,9 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
if not (config.get("SQLLAB_CTAS_NO_LIMIT") and execution_context.select_as_cta): if not (config.get("SQLLAB_CTAS_NO_LIMIT") and execution_context.select_as_cta):
# set LIMIT after template processing # set LIMIT after template processing
db_engine_spec = execution_context.database.db_engine_spec # type: ignore
limits = [ limits = [
execution_context.database.db_engine_spec.get_limit_from_sql( # type: ignore db_engine_spec.get_limit_from_sql(rendered_query),
rendered_query
),
execution_context.limit, execution_context.limit,
] ]
if limits[0] is None or limits[0] > limits[1]: # type: ignore if limits[0] is None or limits[0] > limits[1]: # type: ignore
@@ -2672,11 +2680,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
# Flag for whether or not to expand data # Flag for whether or not to expand data
# (feature that will expand Presto row objects and arrays) # (feature that will expand Presto row objects and arrays)
expand_data: bool = cast( expand_data: bool = execution_context.expand_data
bool,
is_feature_enabled("PRESTO_EXPAND_DATA")
and query_params.get("expand_data"),
)
# Async request. # Async request.
if execution_context.is_run_asynchronous(): if execution_context.is_run_asynchronous():

View File

@@ -63,7 +63,6 @@ QUERY_2 = "SELECT * FROM NO_TABLE"
QUERY_3 = "SELECT * FROM birth_names LIMIT 10" QUERY_3 = "SELECT * FROM birth_names LIMIT 10"
@pytest.mark.sqllab
class TestSqlLab(SupersetTestCase): class TestSqlLab(SupersetTestCase):
"""Testings for Sql Lab""" """Testings for Sql Lab"""