From 76f0408932a34a398969ff7a99fc5ec834be9360 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Thu, 23 Sep 2021 15:18:17 -0700
Subject: [PATCH] feat: handle temporal columns in group bys (#16795)

* feat: handle temporal columns in group bys

* Rebase
---
 superset/connectors/sqla/models.py | 17 ++++++++++++++---
 superset/db_engine_specs/druid.py  | 14 ++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index 2b5c6c85fb..c87b4c2145 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -35,6 +35,7 @@ from typing import (
     Union,
 )
 
+import dateutil.parser
 import pandas as pd
 import sqlalchemy as sa
 import sqlparse
@@ -1416,15 +1417,25 @@ class SqlaTable(Model, BaseDatasource):  # pylint: disable=too-many-public-metho
             )
         return ob
 
-    @staticmethod
     def _get_top_groups(
-        df: pd.DataFrame, dimensions: List[str], groupby_exprs: Dict[str, Any],
+        self, df: pd.DataFrame, dimensions: List[str], groupby_exprs: Dict[str, Any],
     ) -> ColumnElement:
+        column_map = {column.column_name: column for column in self.columns}
         groups = []
         for _unused, row in df.iterrows():
             group = []
             for dimension in dimensions:
-                group.append(groupby_exprs[dimension] == row[dimension])
+                value = row[dimension]
+
+                # Some databases like Druid will return timestamps as strings, but
+                # do not perform automatic casting when comparing these strings to
+                # a timestamp. For cases like this we convert the value from a
+                # string into a timestamp.
+                if column_map[dimension].is_temporal and isinstance(value, str):
+                    dttm = dateutil.parser.parse(value)
+                    value = text(self.db_engine_spec.convert_dttm("TIMESTAMP", dttm))
+
+                group.append(groupby_exprs[dimension] == value)
             groups.append(and_(*group))
 
         return or_(*groups)
diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py
index 61592911ec..0230a8c178 100644
--- a/superset/db_engine_specs/druid.py
+++ b/superset/db_engine_specs/druid.py
@@ -101,3 +101,17 @@ class DruidEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-method
         if tt in (utils.TemporalType.DATETIME, utils.TemporalType.TIMESTAMP):
             return f"""TIME_PARSE('{dttm.isoformat(timespec="seconds")}')"""
         return None
+
+    @classmethod
+    def epoch_to_dttm(cls) -> str:
+        """
+        Convert from number of seconds since the epoch to a timestamp.
+        """
+        return "MILLIS_TO_TIMESTAMP({col} * 1000)"
+
+    @classmethod
+    def epoch_ms_to_dttm(cls) -> str:
+        """
+        Convert from number of milliseconds since the epoch to a timestamp.
+        """
+        return "MILLIS_TO_TIMESTAMP({col})"