Fix time granularity-related issues (apache#4821)

* Fixing time grain * Add tests
wenchma · Apr 18, 2018 · 2ffe5ae · 2ffe5ae
1 parent 9000999
commit 2ffe5ae
Show file tree

Hide file tree

Showing 5 changed files with 135 additions and 110 deletions.
diff --git a/superset/assets/src/explore/stores/visTypes.js b/superset/assets/src/explore/stores/visTypes.js
@@ -95,6 +95,15 @@ export const sections = {
   ],
 };
 
+const timeGrainSqlaAnimationOverrides = {
+  default: null,
+  mapStateToProps: state => ({
+    choices: (state.datasource) ?
+      state.datasource.time_grain_sqla.filter(o => o[0] !== null) :
+      null,
+  }),
+};
+
 export const visTypes = {
   dist_bar: {
     label: t('Distribution - Bar Chart'),
@@ -555,6 +564,7 @@ export const visTypes = {
         description: t("Metric used as a weight for the grid's coloring"),
         validators: [v.nonEmpty],
       },
+      time_grain_sqla: timeGrainSqlaAnimationOverrides,
     },
   },
 
@@ -738,6 +748,7 @@ export const visTypes = {
       size: {
         validators: [],
       },
+      time_grain_sqla: timeGrainSqlaAnimationOverrides,
     },
   },
 

diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
@@ -129,9 +129,10 @@ def get_timestamp_expression(self, time_grain):
                     expr = db_spec.epoch_to_dttm().format(col=expr)
                 elif pdf == 'epoch_ms':
                     expr = db_spec.epoch_ms_to_dttm().format(col=expr)
-            grain = self.table.database.grains_dict().get(time_grain, '{col}')
-            expr = grain.function.format(col=expr)
-        return literal_column(expr, type_=DateTime).label(DTTM_ALIAS)
+            grain = self.table.database.grains_dict().get(time_grain)
+            literal = grain.function if grain else '{col}'
+            literal = expr.format(col=expr)
+        return literal_column(literal, type_=DateTime).label(DTTM_ALIAS)
 
     @classmethod
     def import_obj(cls, i_column):

diff --git a/superset/data/__init__.py b/superset/data/__init__.py
@@ -1301,7 +1301,7 @@ def load_deck_dash():
         "row_limit": 5000,
         "since": None,
         "size": "count",
-        "time_grain_sqla": "Time Column",
+        "time_grain_sqla": None,
         "until": None,
         "viewport": {
             "bearing": -4.952916738791771,
@@ -1359,7 +1359,7 @@ def load_deck_dash():
         },
         "point_radius_fixed": {"type": "fix", "value": 2000},
         "datasource": "5__table",
-        "time_grain_sqla": "Time Column",
+        "time_grain_sqla": None,
         "groupby": [],
     }
     print("Creating Screen Grid slice")
@@ -1408,7 +1408,7 @@ def load_deck_dash():
         "where": "",
         "point_radius_fixed": {"type": "fix", "value": 2000},
         "datasource": "5__table",
-        "time_grain_sqla": "Time Column",
+        "time_grain_sqla": None,
         "groupby": [],
     }
     print("Creating Hex slice")
@@ -1457,7 +1457,7 @@ def load_deck_dash():
         "where": "",
         "point_radius_fixed": {"type": "fix", "value": 2000},
         "datasource": "5__table",
-        "time_grain_sqla": "Time Column",
+        "time_grain_sqla": None,
         "groupby": [],
     }
     print("Creating Grid slice")
@@ -1474,62 +1474,62 @@ def load_deck_dash():
     polygon_tbl = db.session.query(TBL) \
                     .filter_by(table_name='sf_population_polygons').first()
     slice_data = {
-            "datasource": "11__table",
-            "viz_type": "deck_polygon",
-            "slice_id": 41,
-            "granularity_sqla": None,
-            "time_grain_sqla": None,
-            "since": None,
-            "until": None,
-            "line_column": "contour",
-            "line_type": "json",
-            "mapbox_style": "mapbox://styles/mapbox/light-v9",
-            "viewport": {
-                "longitude": -122.43388541747726,
-                "latitude": 37.752020331384834,
-                "zoom": 11.133995608594631,
-                "bearing": 37.89506450385642,
-                "pitch": 60,
-                "width": 667,
-                "height": 906,
-                "altitude": 1.5,
-                "maxZoom": 20,
-                "minZoom": 0,
-                "maxPitch": 60,
-                "minPitch": 0,
-                "maxLatitude": 85.05113,
-                "minLatitude": -85.05113
-            },
-            "reverse_long_lat": False,
-            "fill_color_picker": {
-                "r": 3,
-                "g": 65,
-                "b": 73,
-                "a": 1
-            },
-            "stroke_color_picker": {
-                "r": 0,
-                "g": 122,
-                "b": 135,
-                "a": 1
-            },
-            "filled": True,
-            "stroked": False,
-            "extruded": True,
-            "point_radius_scale": 100,
-            "js_columns": [
-                "population",
-                "area"
-            ],
-            "js_datapoint_mutator": "(d) => {\n    d.elevation = d.extraProps.population/d.extraProps.area/10\n \
-             d.fillColor = [d.extraProps.population/d.extraProps.area/60,140,0]\n \
-             return d;\n}",
-            "js_tooltip": "",
-            "js_onclick_href": "",
-            "where": "",
-            "having": "",
-            "filters": []
-        }
+        "datasource": "11__table",
+        "viz_type": "deck_polygon",
+        "slice_id": 41,
+        "granularity_sqla": None,
+        "time_grain_sqla": None,
+        "since": None,
+        "until": None,
+        "line_column": "contour",
+        "line_type": "json",
+        "mapbox_style": "mapbox://styles/mapbox/light-v9",
+        "viewport": {
+            "longitude": -122.43388541747726,
+            "latitude": 37.752020331384834,
+            "zoom": 11.133995608594631,
+            "bearing": 37.89506450385642,
+            "pitch": 60,
+            "width": 667,
+            "height": 906,
+            "altitude": 1.5,
+            "maxZoom": 20,
+            "minZoom": 0,
+            "maxPitch": 60,
+            "minPitch": 0,
+            "maxLatitude": 85.05113,
+            "minLatitude": -85.05113
+        },
+        "reverse_long_lat": False,
+        "fill_color_picker": {
+            "r": 3,
+            "g": 65,
+            "b": 73,
+            "a": 1
+        },
+        "stroke_color_picker": {
+            "r": 0,
+            "g": 122,
+            "b": 135,
+            "a": 1
+        },
+        "filled": True,
+        "stroked": False,
+        "extruded": True,
+        "point_radius_scale": 100,
+        "js_columns": [
+            "population",
+            "area"
+        ],
+        "js_datapoint_mutator": "(d) => {\n    d.elevation = d.extraProps.population/d.extraProps.area/10\n \
+         d.fillColor = [d.extraProps.population/d.extraProps.area/60,140,0]\n \
+         return d;\n}",
+        "js_tooltip": "",
+        "js_onclick_href": "",
+        "where": "",
+        "having": "",
+        "filters": []
+    }
 
     print("Creating Polygon slice")
     slc = Slice(
@@ -1543,52 +1543,52 @@ def load_deck_dash():
     slices.append(slc)
 
     slice_data = {
-            "datasource": "10__table",
-            "viz_type": "deck_arc",
-            "slice_id": 42,
-            "granularity_sqla": "dttm",
-            "time_grain_sqla": "Time Column",
-            "since": None,
-            "until": None,
-            "start_spatial": {
-                "type": "latlong",
-                "latCol": "LATITUDE",
-                "lonCol": "LONGITUDE"
-            },
-            "end_spatial": {
-                "type": "latlong",
-                "latCol": "LATITUDE_DEST",
-                "lonCol": "LONGITUDE_DEST"
-            },
-            "row_limit": 5000,
-            "mapbox_style": "mapbox://styles/mapbox/light-v9",
-            "viewport": {
-                "altitude": 1.5,
-                "bearing": 8.546256357301871,
-                "height": 642,
-                "latitude": 44.596651438714254,
-                "longitude": -91.84340711201104,
-                "maxLatitude": 85.05113,
-                "maxPitch": 60,
-                "maxZoom": 20,
-                "minLatitude": -85.05113,
-                "minPitch": 0,
-                "minZoom": 0,
-                "pitch": 60,
-                "width": 997,
-                "zoom": 2.929837070560775
-            },
-            "color_picker": {
-                "r": 0,
-                "g": 122,
-                "b": 135,
-                "a": 1
-            },
-            "stroke_width": 1,
-            "where": "",
-            "having": "",
-            "filters": []
-        }
+        "datasource": "10__table",
+        "viz_type": "deck_arc",
+        "slice_id": 42,
+        "granularity_sqla": "dttm",
+        "time_grain_sqla": "Time Column",
+        "since": None,
+        "until": None,
+        "start_spatial": {
+            "type": "latlong",
+            "latCol": "LATITUDE",
+            "lonCol": "LONGITUDE"
+        },
+        "end_spatial": {
+            "type": "latlong",
+            "latCol": "LATITUDE_DEST",
+            "lonCol": "LONGITUDE_DEST"
+        },
+        "row_limit": 5000,
+        "mapbox_style": "mapbox://styles/mapbox/light-v9",
+        "viewport": {
+            "altitude": 1.5,
+            "bearing": 8.546256357301871,
+            "height": 642,
+            "latitude": 44.596651438714254,
+            "longitude": -91.84340711201104,
+            "maxLatitude": 85.05113,
+            "maxPitch": 60,
+            "maxZoom": 20,
+            "minLatitude": -85.05113,
+            "minPitch": 0,
+            "minZoom": 0,
+            "pitch": 60,
+            "width": 997,
+            "zoom": 2.929837070560775
+        },
+        "color_picker": {
+            "r": 0,
+            "g": 122,
+            "b": 135,
+            "a": 1
+        },
+        "stroke_width": 1,
+        "where": "",
+        "having": "",
+        "filters": []
+    }
 
     print("Creating Arc slice")
     slc = Slice(

diff --git a/superset/models/core.py b/superset/models/core.py
@@ -786,7 +786,12 @@ def grains(self):
         return self.db_engine_spec.time_grains
 
     def grains_dict(self):
-        return {grain.duration: grain for grain in self.grains()}
+        """Allowing to lookup grain by either label or duration
+
+        For backward compatibility"""
+        d = {grain.duration: grain for grain in self.grains()}
+        d.update({grain.label: grain for grain in self.grains()})
+        return d
 
     def get_extra(self):
         extra = {}

diff --git a/tests/model_tests.py b/tests/model_tests.py
@@ -97,3 +97,11 @@ def test_select_star(self):
         FROM bart_lines
         LIMIT 100""".format(**locals()))
         assert sql.startswith(expected)
+
+    def test_grains_dict(self):
+        uri = 'mysql://root@localhost'
+        database = Database(sqlalchemy_uri=uri)
+        d = database.grains_dict()
+        self.assertEquals(d.get('day').function, 'DATE({col})')
+        self.assertEquals(d.get('P1D').function, 'DATE({col})')
+        self.assertEquals(d.get('Time Column').function, '{col}')