From 90dbe8d340f227e3ffe374984a202516d779d5d8 Mon Sep 17 00:00:00 2001
From: Jesse Yang <jesse.yang@airbnb.com>
Date: Fri, 1 Apr 2022 16:44:41 -0700
Subject: [PATCH] perf: speed up db migration for deprecating
 time_range_endpoints (#19495)

---
 ...0b36b94_rm_time_range_endpoints_from_qc.py | 22 +++++++++----------
 ...9d86e695_deprecate_time_range_endpoints.py |  4 ++--
 ...dd_saved_query_foreign_key_to_tab_state.py |  1 +
 .../versions/e866bd2d4976_smaller_grid.py     |  1 +
 4 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/superset/migrations/versions/2ed890b36b94_rm_time_range_endpoints_from_qc.py b/superset/migrations/versions/2ed890b36b94_rm_time_range_endpoints_from_qc.py
index 34b833dabd20..42d73bc33d33 100644
--- a/superset/migrations/versions/2ed890b36b94_rm_time_range_endpoints_from_qc.py
+++ b/superset/migrations/versions/2ed890b36b94_rm_time_range_endpoints_from_qc.py
@@ -30,7 +30,6 @@
 
 import sqlalchemy as sa
 from alembic import op
-from sqlalchemy.dialects import postgresql
 from sqlalchemy.ext.declarative import declarative_base
 
 from superset import db
@@ -47,16 +46,17 @@ class Slice(Base):
 def upgrade():
     bind = op.get_bind()
     session = db.Session(bind=bind)
-    for slc in session.query(Slice):
-        if slc.query_context:
-            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)
+    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()
diff --git a/superset/migrations/versions/ab9a9d86e695_deprecate_time_range_endpoints.py b/superset/migrations/versions/ab9a9d86e695_deprecate_time_range_endpoints.py
index 3f149d3ff0c2..2bd65996140d 100644
--- a/superset/migrations/versions/ab9a9d86e695_deprecate_time_range_endpoints.py
+++ b/superset/migrations/versions/ab9a9d86e695_deprecate_time_range_endpoints.py
@@ -46,8 +46,8 @@ def upgrade():
     bind = op.get_bind()
     session = db.Session(bind=bind)
 
-    for slc in session.query(Slice):
-        params = json.loads(slc.params or "{}")
+    for slc in session.query(Slice).filter(Slice.params.like("%time_range_endpoints%")):
+        params = json.loads(slc.params)
         params.pop("time_range_endpoints", None)
         slc.params = json.dumps(params)
 
diff --git a/superset/migrations/versions/c53bae8f08dd_add_saved_query_foreign_key_to_tab_state.py b/superset/migrations/versions/c53bae8f08dd_add_saved_query_foreign_key_to_tab_state.py
index a70b95af0ad4..2a2d66a1fd19 100644
--- a/superset/migrations/versions/c53bae8f08dd_add_saved_query_foreign_key_to_tab_state.py
+++ b/superset/migrations/versions/c53bae8f08dd_add_saved_query_foreign_key_to_tab_state.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """add_saved_query_foreign_key_to_tab_state
+
 Revision ID: c53bae8f08dd
 Revises: bb38f40aa3ff
 Create Date: 2021-12-15 15:05:21.845777
diff --git a/superset/migrations/versions/e866bd2d4976_smaller_grid.py b/superset/migrations/versions/e866bd2d4976_smaller_grid.py
index b674fd603073..286be8a5fc9c 100644
--- a/superset/migrations/versions/e866bd2d4976_smaller_grid.py
+++ b/superset/migrations/versions/e866bd2d4976_smaller_grid.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """smaller_grid
+
 Revision ID: e866bd2d4976
 Revises: 21e88bc06c02
 Create Date: 2018-02-13 08:07:40.766277