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

superset/utils/sqllab_execution_context.py

@@ -105,7 +105,7 @@ class SqlJsonExecutionContext:  # pylint: disable=too-many-instance-attributes
             limit = 0
         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:
             return g.user.get_id() if g.user else None
         except RuntimeError:
@@ -135,7 +135,7 @@ class SqlJsonExecutionContext:  # pylint: disable=too-many-instance-attributes
         pass
 
     def create_query(self) -> Query:
-        # pylint: disable=C0301
+        # pylint: disable=line-too-long
         start_time = now_as_float()
         if self.select_as_cta:
             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
     target_schema_name: Optional[str]
     target_table_name: str

superset/views/core.py

@@ -2556,7 +2556,7 @@ class Superset(BaseSupersetView):  # pylint: disable=too-many-public-methods
             "user_agent": cast(Optional[str], request.headers.get("USER_AGENT"))
         }
         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
     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
         self,
         execution_context: SqlJsonExecutionContext,
-        query_params: Dict[str, Any],
         log_params: Optional[Dict[str, Any]] = None,
     ) -> FlaskResponse:
         """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))
             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(
             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):
             # set LIMIT after template processing
+            db_engine_spec = execution_context.database.db_engine_spec  # type: ignore
             limits = [
-                execution_context.database.db_engine_spec.get_limit_from_sql(  # type: ignore
-                    rendered_query
-                ),
+                db_engine_spec.get_limit_from_sql(rendered_query),
                 execution_context.limit,
             ]
             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
         # (feature that will expand Presto row objects and arrays)
-        expand_data: bool = cast(
-            bool,
-            is_feature_enabled("PRESTO_EXPAND_DATA")
-            and query_params.get("expand_data"),
-        )
+        expand_data: bool = execution_context.expand_data
 
         # Async request.
         if execution_context.is_run_asynchronous():

tests/integration_tests/sqllab_tests.py

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