diff --git a/superset/common/query_context_processor.py b/superset/common/query_context_processor.py index 98aaecebd9bfd..01259ede1d8a5 100644 --- a/superset/common/query_context_processor.py +++ b/superset/common/query_context_processor.py @@ -266,7 +266,8 @@ def _get_timestamp_format( # Query datasource didn't support `get_column` and hasattr(datasource, "get_column") and (col := datasource.get_column(label)) - and col.is_dttm + # todo(hugh) standardize column object in Query datasource + and (col.get("is_dttm") if isinstance(col, dict) else col.is_dttm) ) dttm_cols = [ DateColumn( diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py index 408bc708dfb81..f75973ad1794f 100644 --- a/superset/models/sql_lab.py +++ b/superset/models/sql_lab.py @@ -51,7 +51,6 @@ ) from superset.sql_parse import CtasMethod, ParsedQuery, Table from superset.sqllab.limiting_factor import LimitingFactor -from superset.superset_typing import ResultSetColumnType from superset.utils.core import GenericDataType, QueryStatus, user_label if TYPE_CHECKING: @@ -183,7 +182,7 @@ def sql_tables(self) -> List[Table]: return list(ParsedQuery(self.sql).tables) @property - def columns(self) -> List[ResultSetColumnType]: + def columns(self) -> List[Dict[str, Any]]: bool_types = ("BOOL",) num_types = ( "DOUBLE", @@ -217,7 +216,7 @@ def columns(self) -> List[ResultSetColumnType]: computed_column["column_name"] = col.get("name") computed_column["groupby"] = True columns.append(computed_column) - return columns # type: ignore + return columns @property def data(self) -> Dict[str, Any]: @@ -288,7 +287,7 @@ def offset(self) -> int: def main_dttm_col(self) -> Optional[str]: for col in self.columns: if col.get("is_dttm"): - return col.get("column_name") # type: ignore + return col.get("column_name") return None @property @@ -332,6 +331,14 @@ def tracking_url(self) -> Optional[str]: def tracking_url(self, value: str) -> None: self.tracking_url_raw = value + def get_column(self, column_name: Optional[str]) -> Optional[Dict[str, Any]]: + if not column_name: + return None + for col in self.columns: + if col.get("column_name") == column_name: + return col + return None + class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin, ImportExportMixin): """ORM model for SQL query"""