feat(sqllab): TRINO_EXPAND_ROWS: expand columns from ROWs (apache#25809)

superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx

@@ -202,7 +202,7 @@ const ExtraOptions = ({
                 />
               </div>
             </StyledInputContainer>
-            <StyledInputContainer>
+            <StyledInputContainer css={no_margin_bottom}>
               <div className="input-container">
                 <IndeterminateCheckbox
                   id="disable_data_preview"
@@ -220,6 +220,22 @@ const ExtraOptions = ({
                 />
               </div>
             </StyledInputContainer>
+            <StyledInputContainer>
+              <div className="input-container">
+                <IndeterminateCheckbox
+                  id="expand_rows"
+                  indeterminate={false}
+                  checked={!!extraJson?.schema_options?.expand_rows}
+                  onChange={onExtraInputChange}
+                  labelText={t('Enable row expansion in schemas')}
+                />
+                <InfoTooltip
+                  tooltip={t(
+                    'For Trino, describe full schemas of nested ROW types, expanding them with dotted paths',
+                  )}
+                />
+              </div>
+            </StyledInputContainer>
           </StyledExpandableForm>
         </StyledInputContainer>
       </Collapse.Panel>

superset-frontend/src/features/databases/DatabaseModal/index.test.tsx

@@ -674,7 +674,7 @@ describe('DatabaseModal', () => {
       const exposeInSQLLabCheckbox = screen.getByRole('checkbox', {
         name: /expose database in sql lab/i,
       });
-      // This is both the checkbox and it's respective SVG
+      // This is both the checkbox and its respective SVG
       // const exposeInSQLLabCheckboxSVG = checkboxOffSVGs[0].parentElement;
       const exposeInSQLLabText = screen.getByText(
         /expose database in sql lab/i,
@@ -721,6 +721,13 @@ describe('DatabaseModal', () => {
         /Disable SQL Lab data preview queries/i,
       );
 
+      const enableRowExpansionCheckbox = screen.getByRole('checkbox', {
+        name: /enable row expansion in schemas/i,
+      });
+      const enableRowExpansionText = screen.getByText(
+        /enable row expansion in schemas/i,
+      );
+
       // ---------- Assertions ----------
       const visibleComponents = [
         closeButton,
@@ -737,13 +744,15 @@ describe('DatabaseModal', () => {
         checkboxOffSVGs[2],
         checkboxOffSVGs[3],
         checkboxOffSVGs[4],
+        checkboxOffSVGs[5],
         tooltipIcons[0],
         tooltipIcons[1],
         tooltipIcons[2],
         tooltipIcons[3],
         tooltipIcons[4],
         tooltipIcons[5],
         tooltipIcons[6],
+        tooltipIcons[7],
         exposeInSQLLabText,
         allowCTASText,
         allowCVASText,
@@ -754,6 +763,7 @@ describe('DatabaseModal', () => {
         enableQueryCostEstimationText,
         allowDbExplorationText,
         disableSQLLabDataPreviewQueriesText,
+        enableRowExpansionText,
       ];
       // These components exist in the DOM but are not visible
       const invisibleComponents = [
@@ -764,15 +774,16 @@ describe('DatabaseModal', () => {
         enableQueryCostEstimationCheckbox,
         allowDbExplorationCheckbox,
         disableSQLLabDataPreviewQueriesCheckbox,
+        enableRowExpansionCheckbox,
       ];
       visibleComponents.forEach(component => {
         expect(component).toBeVisible();
       });
       invisibleComponents.forEach(component => {
         expect(component).not.toBeVisible();
       });
-      expect(checkboxOffSVGs).toHaveLength(5);
-      expect(tooltipIcons).toHaveLength(7);
+      expect(checkboxOffSVGs).toHaveLength(6);
+      expect(tooltipIcons).toHaveLength(8);
     });
 
     test('renders the "Advanced" - PERFORMANCE tab correctly', async () => {

superset-frontend/src/features/databases/DatabaseModal/index.tsx

@@ -307,6 +307,18 @@ export function dbReducer(
           }),
         };
       }
+      if (action.payload.name === 'expand_rows') {
+        return {
+          ...trimmedState,
+          extra: JSON.stringify({
+            ...extraJson,
+            schema_options: {
+              ...extraJson?.schema_options,
+              [action.payload.name]: !!action.payload.value,
+            },
+          }),
+        };
+      }
       return {
         ...trimmedState,
         extra: JSON.stringify({

superset-frontend/src/features/databases/types.ts

@@ -226,5 +226,8 @@ export interface ExtraJson {
     table_cache_timeout?: number; // in Performance
   }; // No field, holds schema and table timeout
   schemas_allowed_for_file_upload?: string[]; // in Security
+  schema_options?: {
+    expand_rows?: boolean;
+  };
   version?: string;
 }

superset/db_engine_specs/base.py

@@ -51,7 +51,7 @@
 from sqlalchemy.engine.url import URL
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import Session
-from sqlalchemy.sql import quoted_name, text
+from sqlalchemy.sql import literal_column, quoted_name, text
 from sqlalchemy.sql.expression import ColumnClause, Select, TextAsFrom, TextClause
 from sqlalchemy.types import TypeEngine
 from sqlparse.tokens import CTE
@@ -1322,15 +1322,21 @@ def get_table_comment(
         return comment
 
     @classmethod
-    def get_columns(
-        cls, inspector: Inspector, table_name: str, schema: str | None
+    def get_columns(  # pylint: disable=unused-argument
+        cls,
+        inspector: Inspector,
+        table_name: str,
+        schema: str | None,
+        options: dict[str, Any] | None = None,
     ) -> list[ResultSetColumnType]:
         """
         Get all columns from a given schema and table
 
         :param inspector: SqlAlchemy Inspector instance
         :param table_name: Table name
         :param schema: Schema name. If omitted, uses default schema for database
+        :param options: Extra options to customise the display of columns in
+                        some databases
         :return: All columns in table
         """
         return convert_inspector_columns(
@@ -1382,7 +1388,12 @@ def where_latest_partition(  # pylint: disable=too-many-arguments,unused-argumen
 
     @classmethod
     def _get_fields(cls, cols: list[ResultSetColumnType]) -> list[Any]:
-        return [column(c["column_name"]) for c in cols]
+        return [
+            literal_column(query_as)
+            if (query_as := c.get("query_as"))
+            else column(c["column_name"])
+            for c in cols
+        ]
 
     @classmethod
     def select_star(  # pylint: disable=too-many-arguments,too-many-locals

superset/db_engine_specs/druid.py

@@ -23,14 +23,12 @@
 from typing import Any, TYPE_CHECKING
 
 from sqlalchemy import types
-from sqlalchemy.engine.reflection import Inspector
 
 from superset import is_feature_enabled
 from superset.constants import TimeGrain
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.exceptions import SupersetDBAPIConnectionError
 from superset.exceptions import SupersetException
-from superset.superset_typing import ResultSetColumnType
 from superset.utils import core as utils
 
 if TYPE_CHECKING:
@@ -130,15 +128,6 @@ def epoch_ms_to_dttm(cls) -> str:
         """
         return "MILLIS_TO_TIMESTAMP({col})"
 
-    @classmethod
-    def get_columns(
-        cls, inspector: Inspector, table_name: str, schema: str | None
-    ) -> list[ResultSetColumnType]:
-        """
-        Update the Druid type map.
-        """
-        return super().get_columns(inspector, table_name, schema)
-
     @classmethod
     def get_dbapi_exception_mapping(cls) -> dict[type[Exception], type[Exception]]:
         # pylint: disable=import-outside-toplevel

superset/db_engine_specs/hive.py

@@ -410,9 +410,13 @@ def handle_cursor(  # pylint: disable=too-many-locals
 
     @classmethod
     def get_columns(
-        cls, inspector: Inspector, table_name: str, schema: str | None
+        cls,
+        inspector: Inspector,
+        table_name: str,
+        schema: str | None,
+        options: dict[str, Any] | None = None,
     ) -> list[ResultSetColumnType]:
-        return BaseEngineSpec.get_columns(inspector, table_name, schema)
+        return BaseEngineSpec.get_columns(inspector, table_name, schema, options)
 
     @classmethod
     def where_latest_partition(  # pylint: disable=too-many-arguments

superset/db_engine_specs/presto.py

@@ -981,14 +981,19 @@ def _show_columns(
 
     @classmethod
     def get_columns(
-        cls, inspector: Inspector, table_name: str, schema: str | None
+        cls,
+        inspector: Inspector,
+        table_name: str,
+        schema: str | None,
+        options: dict[str, Any] | None = None,
     ) -> list[ResultSetColumnType]:
         """
         Get columns from a Presto data source. This includes handling row and
         array data types
         :param inspector: object that performs database schema inspection
         :param table_name: table name
         :param schema: schema name
+        :param options: Extra configuration options, not used by this backend
         :return: a list of results that contain column info
                 (i.e. column name and data type)
         """

superset/db_engine_specs/trino.py

@@ -24,15 +24,18 @@
 
 import simplejson as json
 from flask import current_app
+from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.engine.url import URL
 from sqlalchemy.orm import Session
+from trino.sqlalchemy import datatype
 
 from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY, USER_AGENT
 from superset.databases.utils import make_url_safe
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.exceptions import SupersetDBAPIConnectionError
 from superset.db_engine_specs.presto import PrestoBaseEngineSpec
 from superset.models.sql_lab import Query
+from superset.superset_typing import ResultSetColumnType
 from superset.utils import core as utils
 
 if TYPE_CHECKING:
@@ -331,3 +334,62 @@ def get_dbapi_exception_mapping(cls) -> dict[type[Exception], type[Exception]]:
         return {
             requests_exceptions.ConnectionError: SupersetDBAPIConnectionError,
         }
+
+    @classmethod
+    def _expand_columns(cls, col: ResultSetColumnType) -> list[ResultSetColumnType]:
+        """
+        Expand the given column out to one or more columns by analysing their types,
+        descending into ROWS and expanding out their inner fields recursively.
+
+        We can only navigate named fields in ROWs in this way, so we can't expand out
+        MAP or ARRAY types, nor fields in ROWs which have no name (in fact the trino
+        library doesn't correctly parse unnamed fields in ROWs). We won't be able to
+        expand ROWs which are nested underneath any of those types, either.
+
+        Expanded columns are named foo.bar.baz and we provide a query_as property to
+        instruct the base engine spec how to correctly query them: instead of quoting
+        the whole string they have to be quoted like "foo"."bar"."baz" and we then
+        alias them to the full dotted string for ease of reference.
+        """
+        cols = [col]
+        col_type = col.get("type")
+
+        if not isinstance(col_type, datatype.ROW):
+            return cols
+
+        for inner_name, inner_type in col_type.attr_types:
+            outer_name = col["name"]
+            name = ".".join([outer_name, inner_name])
+            query_name = ".".join([f'"{piece}"' for piece in name.split(".")])
+            column_spec = cls.get_column_spec(str(inner_type))
+            is_dttm = column_spec.is_dttm if column_spec else False
+
+            inner_col = ResultSetColumnType(
+                name=name,
+                column_name=name,
+                type=inner_type,
+                is_dttm=is_dttm,
+                query_as=f'{query_name} AS "{name}"',
+            )
+            cols.extend(cls._expand_columns(inner_col))
+
+        return cols
+
+    @classmethod
+    def get_columns(
+        cls,
+        inspector: Inspector,
+        table_name: str,
+        schema: str | None,
+        options: dict[str, Any] | None = None,
+    ) -> list[ResultSetColumnType]:
+        """
+        If the "expand_rows" feature is enabled on the database via
+        "schema_options", expand the schema definition out to show all
+        subfields of nested ROWs as their appropriate dotted paths.
+        """
+        base_cols = super().get_columns(inspector, table_name, schema, options)
+        if not (options or {}).get("expand_rows"):
+            return base_cols
+
+        return [col for base_col in base_cols for col in cls._expand_columns(base_col)]

superset/models/core.py

@@ -237,6 +237,11 @@ def disable_data_preview(self) -> bool:
         # this will prevent any 'trash value' strings from going through
         return self.get_extra().get("disable_data_preview", False) is True
 
+    @property
+    def schema_options(self) -> dict[str, Any]:
+        """Additional schema display config for engines with complex schemas"""
+        return self.get_extra().get("schema_options", {})
+
     @property
     def data(self) -> dict[str, Any]:
         return {
@@ -248,6 +253,7 @@ def data(self) -> dict[str, Any]:
             "allows_cost_estimate": self.allows_cost_estimate,
             "allows_virtual_table_explore": self.allows_virtual_table_explore,
             "explore_database_id": self.explore_database_id,
+            "schema_options": self.schema_options,
             "parameters": self.parameters,
             "disable_data_preview": self.disable_data_preview,
             "parameters_schema": self.parameters_schema,
@@ -838,7 +844,9 @@ def get_columns(
         self, table_name: str, schema: str | None = None
     ) -> list[ResultSetColumnType]:
         with self.get_inspector_with_context() as inspector:
-            return self.db_engine_spec.get_columns(inspector, table_name, schema)
+            return self.db_engine_spec.get_columns(
+                inspector, table_name, schema, self.schema_options
+            )
 
     def get_metrics(
         self,

superset/superset_typing.py

@@ -84,6 +84,8 @@ class ResultSetColumnType(TypedDict):
     scale: NotRequired[Any]
     max_length: NotRequired[Any]
 
+    query_as: NotRequired[Any]
+
 
 CacheConfig = dict[str, Any]
 DbapiDescriptionRow = tuple[