use dynamic time_grains for schema (apache#14009)

preset-io · Apr 8, 2021 · 667eb83 · 667eb83
1 parent c60a93d
commit 667eb83
Showing 1 changed file with 18 additions and 22 deletions.
diff --git a/superset/charts/schemas.py b/superset/charts/schemas.py
@@ -21,7 +21,9 @@
 from marshmallow.validate import Length, Range
 from marshmallow_enum import EnumField
 
+from superset import app
 from superset.common.query_context import QueryContext
+from superset.db_engine_specs.base import builtin_time_grains
 from superset.utils import schema as utils
 from superset.utils.core import (
     AnnotationType,
@@ -33,6 +35,8 @@
     TimeRangeEndpoint,
 )
 
+config = app.config
+
 #
 # RISON/JSON schemas for query parameters
 #
@@ -126,26 +130,6 @@
 }
 
 
-TIME_GRAINS = (
-    "PT1S",
-    "PT1M",
-    "PT5M",
-    "PT10M",
-    "PT15M",
-    "PT0.5H",
-    "PT1H",
-    "P1D",
-    "P1W",
-    "P1M",
-    "P0.25Y",
-    "P1Y",
-    "1969-12-28T00:00:00Z/P1W",  # Week starting Sunday
-    "1969-12-29T00:00:00Z/P1W",  # Week starting Monday
-    "P1W/1970-01-03T00:00:00Z",  # Week ending Saturday
-    "P1W/1970-01-04T00:00:00Z",  # Week ending Sunday
-)
-
-
 class ChartEntityResponseSchema(Schema):
     """
     Schema for a chart object
@@ -498,7 +482,13 @@ class ChartDataProphetOptionsSchema(ChartDataPostProcessingOperationOptionsSchem
         description="Time grain used to specify time period increments in prediction. "
         "Supports [ISO 8601](https://en.wikipedia.org/wiki/ISO_8601#Durations) "
         "durations.",
-        validate=validate.OneOf(choices=TIME_GRAINS),
+        validate=validate.OneOf(
+            choices=[
+                i
+                for i in {**builtin_time_grains, **config["TIME_GRAIN_ADDONS"]}.keys()
+                if i
+            ]
+        ),
         example="P1D",
         required=True,
     )
@@ -796,7 +786,13 @@ class ChartDataExtrasSchema(Schema):
         description="To what level of granularity should the temporal column be "
         "aggregated. Supports "
         "[ISO 8601](https://en.wikipedia.org/wiki/ISO_8601#Durations) durations.",
-        validate=validate.OneOf(choices=TIME_GRAINS),
+        validate=validate.OneOf(
+            choices=[
+                i
+                for i in {**builtin_time_grains, **config["TIME_GRAIN_ADDONS"]}.keys()
+                if i
+            ]
+        ),
         example="P1D",
         allow_none=True,
     )