From 694ae6f90e76b864d5d721c4223d36e01f7fd4cf Mon Sep 17 00:00:00 2001
From: John Bodley <4567245+john-bodley@users.noreply.github.com>
Date: Wed, 3 Mar 2021 13:22:07 +1300
Subject: [PATCH] fix(query-object): extra time-range-endpoints (#13331)

Co-authored-by: John Bodley <john.bodley@airbnb.com>
---
 superset/charts/schemas.py      | 9 ++-------
 superset/common/query_object.py | 6 ++++--
 tests/fixtures/query_context.py | 7 +++++--
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/superset/charts/schemas.py b/superset/charts/schemas.py
index 185e2aad8..61f41f66a 100644
--- a/superset/charts/schemas.py
+++ b/superset/charts/schemas.py
@@ -30,6 +30,7 @@ from superset.utils.core import (
     FilterOperator,
     PostProcessingBoxplotWhiskerType,
     PostProcessingContributionOrientation,
+    TimeRangeEndpoint,
 )
 
 #
@@ -769,13 +770,7 @@ class ChartDataFilterSchema(Schema):
 
 class ChartDataExtrasSchema(Schema):
 
-    time_range_endpoints = fields.List(
-        fields.String(
-            validate=validate.OneOf(choices=("unknown", "inclusive", "exclusive")),
-            description="A list with two values, stating if start/end should be "
-            "inclusive/exclusive.",
-        )
-    )
+    time_range_endpoints = fields.List(EnumField(TimeRangeEndpoint, by_value=True))
     relative_start = fields.String(
         description="Start time for relative time deltas. "
         'Default: `config["DEFAULT_RELATIVE_START_TIME"]`',
diff --git a/superset/common/query_object.py b/superset/common/query_object.py
index 69baacaa2..4b5f6aaa1 100644
--- a/superset/common/query_object.py
+++ b/superset/common/query_object.py
@@ -181,8 +181,10 @@ class QueryObject:
         self.order_desc = order_desc
         self.extras = extras
 
-        if config["SIP_15_ENABLED"] and "time_range_endpoints" not in self.extras:
-            self.extras["time_range_endpoints"] = get_time_range_endpoints(form_data={})
+        if config["SIP_15_ENABLED"]:
+            self.extras["time_range_endpoints"] = get_time_range_endpoints(
+                form_data=self.extras
+            )
 
         self.columns = columns
         self.groupby = groupby or []
diff --git a/tests/fixtures/query_context.py b/tests/fixtures/query_context.py
index 229549cfe..0bacb671e 100644
--- a/tests/fixtures/query_context.py
+++ b/tests/fixtures/query_context.py
@@ -17,13 +17,16 @@
 import copy
 from typing import Any, Dict, List
 
-from superset.utils.core import AnnotationType, DTTM_ALIAS
+from superset.utils.core import AnnotationType, DTTM_ALIAS, TimeRangeEndpoint
 from tests.base_tests import get_table_by_name
 
 query_birth_names = {
     "extras": {
         "where": "",
-        "time_range_endpoints": ["inclusive", "exclusive"],
+        "time_range_endpoints": (
+            TimeRangeEndpoint.INCLUSIVE,
+            TimeRangeEndpoint.EXCLUSIVE,
+        ),
         "time_grain_sqla": "P1D",
     },
     "groupby": ["name"],