diff --git a/superset/assets/src/SqlLab/actions/sqlLab.js b/superset/assets/src/SqlLab/actions/sqlLab.js index e9785f850da..a540cec28af 100644 --- a/superset/assets/src/SqlLab/actions/sqlLab.js +++ b/superset/assets/src/SqlLab/actions/sqlLab.js @@ -262,6 +262,7 @@ export function runQuery(query) { select_as_cta: query.ctas, templateParams: query.templateParams, queryLimit: query.queryLimit, + expand_data: true, }; return SupersetClient.post({ diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 9050f06d95b..575a10d4d91 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -148,6 +148,7 @@ def get_sql_results( store_results=False, user_name=None, start_time=None, + expand_data=False, ): """Executes the sql query returns the results.""" with session_scope(not ctask.request.called_directly) as session: @@ -162,6 +163,7 @@ def get_sql_results( user_name, session=session, start_time=start_time, + expand_data=expand_data, ) except Exception as e: logging.exception(f"Query {query_id}: {e}") @@ -264,6 +266,7 @@ def _serialize_and_expand_data( cdf: SupersetDataFrame, db_engine_spec: BaseEngineSpec, use_msgpack: Optional[bool] = False, + expand_data: bool = False, ) -> Tuple[Union[bytes, str], list, list, list]: selected_columns: list = cdf.columns or [] expanded_columns: list @@ -282,9 +285,13 @@ def _serialize_and_expand_data( all_columns, expanded_columns = (selected_columns, []) else: data = cdf.data or [] - all_columns, data, expanded_columns = db_engine_spec.expand_data( - selected_columns, data - ) + if expand_data: + all_columns, data, expanded_columns = db_engine_spec.expand_data( + selected_columns, data + ) + else: + all_columns = selected_columns + expanded_columns = [] return (data, selected_columns, all_columns, expanded_columns) @@ -298,6 +305,7 @@ def execute_sql_statements( user_name=None, session=None, start_time=None, + expand_data=False, ): """Executes the sql query returns the results.""" if store_results and start_time: @@ -371,7 +379,7 @@ def execute_sql_statements( query.end_time = now_as_float() data, selected_columns, all_columns, expanded_columns = _serialize_and_expand_data( - cdf, db_engine_spec, store_results and results_backend_use_msgpack + cdf, db_engine_spec, store_results and results_backend_use_msgpack, expand_data ) payload.update( diff --git a/superset/views/core.py b/superset/views/core.py index 01c0cdc706e..944c743278a 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2581,7 +2581,7 @@ class Superset(BaseSupersetView): return json_error_response(f"{msg}") def _sql_json_async( - self, session: Session, rendered_query: str, query: Query + self, session: Session, rendered_query: str, query: Query, expand_data: bool ) -> str: """ Send SQL JSON query to celery workers @@ -2601,6 +2601,7 @@ class Superset(BaseSupersetView): store_results=not query.select_as_cta, user_name=g.user.username if g.user else None, start_time=now_as_float(), + expand_data=expand_data, ) except Exception as e: logging.exception(f"Query {query.id}: {e}") @@ -2625,7 +2626,7 @@ class Superset(BaseSupersetView): return resp def _sql_json_sync( - self, session: Session, rendered_query: str, query: Query + self, session: Session, rendered_query: str, query: Query, expand_data: bool ) -> str: """ Execute SQL query (sql json) @@ -2644,6 +2645,7 @@ class Superset(BaseSupersetView): rendered_query, return_results=True, user_name=g.user.username if g.user else None, + expand_data=expand_data, ) payload = json.dumps( @@ -2756,11 +2758,17 @@ class Superset(BaseSupersetView): limits = [mydb.db_engine_spec.get_limit_from_sql(rendered_query), limit] query.limit = min(lim for lim in limits if lim is not None) + # Flag for whether or not to expand data + # (feature that will expand Presto row objects and arrays) + expand_data: bool = is_feature_enabled( + "PRESTO_EXPAND_DATA" + ) and request.json.get("expand_data") + # Async request. if async_flag: - return self._sql_json_async(session, rendered_query, query) + return self._sql_json_async(session, rendered_query, query, expand_data) # Sync request. - return self._sql_json_sync(session, rendered_query, query) + return self._sql_json_sync(session, rendered_query, query, expand_data) @has_access @expose("/csv/") diff --git a/tests/celery_tests.py b/tests/celery_tests.py index 742e61904bf..36706948ce2 100644 --- a/tests/celery_tests.py +++ b/tests/celery_tests.py @@ -260,7 +260,7 @@ class CeleryTestCase(SupersetTestCase): db_engine_spec, "expand_data", wraps=db_engine_spec.expand_data ) as expand_data: data, selected_columns, all_columns, expanded_columns = sql_lab._serialize_and_expand_data( - cdf, db_engine_spec, False + cdf, db_engine_spec, False, True ) expand_data.assert_called_once()