fix: Removetime_range_endpoints from query context object pt 2 (#19728)

superset/migrations/versions/2ed890b36b94_rm_time_range_endpoints_from_qc.py

@@ -26,40 +26,9 @@
 revision = "2ed890b36b94"
 down_revision = "58df9d617f14"
 
-import json
-
-import sqlalchemy as sa
-from alembic import op
-from sqlalchemy.ext.declarative import declarative_base
-
-from superset import db
-
-Base = declarative_base()
-
-
-class Slice(Base):
-    __tablename__ = "slices"
-    id = sa.Column(sa.Integer, primary_key=True)
-    query_context = sa.Column(sa.Text)
-
 
 def upgrade():
-    bind = op.get_bind()
-    session = db.Session(bind=bind)
-    for slc in session.query(Slice).filter(
-        Slice.query_context.like("%time_range_endpoints%")
-    ):
-        try:
-            query_context = json.loads(slc.query_context)
-        except json.decoder.JSONDecodeError:
-            continue
-        queries = query_context.get("queries")
-        for query in queries:
-            query.get("extras", {}).pop("time_range_endpoints", None)
-        slc.queries = json.dumps(queries)
-
-    session.commit()
-    session.close()
+    pass
 
 
 def downgrade():

superset/migrations/versions/cecc6bf46990_rm_time_range_endpoints_2.py

@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""rm_time_range_endpoints_2
+
+Revision ID: cecc6bf46990
+Revises: 9d8a8d575284
+Create Date: 2022-04-14 17:21:53.996022
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "cecc6bf46990"
+down_revision = "9d8a8d575284"
+
+import json
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.ext.declarative import declarative_base
+
+from superset import db
+
+Base = declarative_base()
+
+
+class Slice(Base):
+    __tablename__ = "slices"
+    id = sa.Column(sa.Integer, primary_key=True)
+    query_context = sa.Column(sa.Text)
+    slice_name = sa.Column(sa.String(250))
+
+
+def upgrade_slice(slc: Slice):
+    try:
+        query_context = json.loads(slc.query_context)
+    except json.decoder.JSONDecodeError:
+        return
+
+    queries = query_context.get("queries")
+
+    for query in queries:
+        query.get("extras", {}).pop("time_range_endpoints", None)
+
+    slc.query_context = json.dumps(query_context)
+
+
+def upgrade():
+    bind = op.get_bind()
+    session = db.Session(bind=bind)
+    for slc in session.query(Slice).filter(
+        Slice.query_context.like("%time_range_endpoints%")
+    ):
+        upgrade_slice(slc)
+
+    session.commit()
+    session.close()
+
+
+def downgrade():
+    pass

tests/integration_tests/migrations/cecc6bf46990_rm_time_range_endpoints_2__tests.py

@@ -0,0 +1,130 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import json
+
+from superset.migrations.versions.cecc6bf46990_rm_time_range_endpoints_2 import (
+    Slice,
+    upgrade_slice,
+)
+
+sample_query_context = {
+    "datasource": {"id": 27, "type": "table"},
+    "force": False,
+    "queries": [
+        {
+            "time_range": "No filter",
+            "filters": [],
+            "extras": {
+                "time_grain_sqla": "P1D",
+                "time_range_endpoints": ["inclusive", "exclusive"],
+                "having": "",
+                "having_druid": [],
+                "where": "",
+            },
+            "applied_time_extras": {},
+            "columns": ["a", "b"],
+            "orderby": [],
+            "annotation_layers": [],
+            "row_limit": 1000,
+            "timeseries_limit": 0,
+            "order_desc": True,
+            "url_params": {},
+            "custom_params": {},
+            "custom_form_data": {},
+            "post_processing": [],
+        }
+    ],
+    "form_data": {
+        "viz_type": "table",
+        "datasource": "27__table",
+        "slice_id": 545,
+        "url_params": {},
+        "time_grain_sqla": "P1D",
+        "time_range": "No filter",
+        "query_mode": "raw",
+        "groupby": [],
+        "metrics": [],
+        "all_columns": ["a", "b"],
+        "percent_metrics": [],
+        "adhoc_filters": [],
+        "order_by_cols": [],
+        "row_limit": 1000,
+        "server_page_length": 10,
+        "include_time": False,
+        "order_desc": True,
+        "table_timestamp_format": "smart_date",
+        "show_cell_bars": True,
+        "color_pn": True,
+        "extra_form_data": {},
+        "force": False,
+        "result_format": "json",
+        "result_type": "full",
+    },
+    "result_format": "json",
+    "result_type": "full",
+}
+
+
+sample_query_context = {
+    "datasource": {"id": 27, "type": "table"},
+    "force": False,
+    "queries": [
+        {
+            "time_range": "No filter",
+            "filters": [],
+            "extras": {
+                "time_grain_sqla": "P1D",
+                "time_range_endpoints": ["inclusive", "exclusive"],
+                "having": "",
+                "having_druid": [],
+                "where": "",
+            },
+            "applied_time_extras": {},
+            "columns": ["a", "b"],
+            "orderby": [],
+            "annotation_layers": [],
+            "row_limit": 1000,
+            "timeseries_limit": 0,
+            "order_desc": True,
+            "url_params": {},
+            "custom_params": {},
+            "custom_form_data": {},
+            "post_processing": [],
+        }
+    ],
+    "form_data": {},
+    "result_format": "json",
+    "result_type": "full",
+}
+
+
+def test_upgrade():
+    slc = Slice(slice_name="FOO", query_context=json.dumps(sample_query_context))
+
+    upgrade_slice(slc)
+
+    query_context = json.loads(slc.query_context)
+    queries = query_context.get("queries")
+    for q in queries:
+        extras = q.get("extras", {})
+        assert "time_range_endpoints" not in extras
+
+
+def test_upgrade_bad_json():
+    slc = Slice(slice_name="FOO", query_context="abc")
+
+    assert None == upgrade_slice(slc)