refactor sql_json view endpoint: separate concern into ad hod method

apache · Sep 5, 2021 · 93690ef · 93690ef
1 parent be77ad2
commit 93690ef
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 14 deletions.
diff --git a/superset/utils/sqllab_execution_context.py b/superset/utils/sqllab_execution_context.py
@@ -105,7 +105,7 @@ def _get_limit_param(query_params: Dict[str, Any]) -> int:
             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 @@ def _validate_db(self, database: Database) -> None:
         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 @@ def create_query(self) -> Query:
         )
 
 
-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

diff --git a/superset/views/core.py b/superset/views/core.py
@@ -2556,7 +2556,7 @@ def sql_json(self) -> FlaskResponse:
             "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 @@ def is_query_handled(cls, query: Optional[Query]) -> bool:
     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 @@ def sql_json_exec(  # pylint: disable=too-many-statements
             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 @@ def sql_json_exec(  # pylint: disable=too-many-statements
 
         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 @@ def sql_json_exec(  # pylint: disable=too-many-statements
 
         # 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():

diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py
@@ -63,7 +63,6 @@
 QUERY_3 = "SELECT * FROM birth_names LIMIT 10"
 
 
-@pytest.mark.sqllab
 class TestSqlLab(SupersetTestCase):
     """Testings for Sql Lab"""