chore: bump majors on celery and Flask (#19168)

* chore: bump celery, Flask, flask-jwt-extended, pyJWT

* fix pyJWT breaking change

* fix pyJWT breaking change 2

* test

* fixed test

* fixed test

* fixed test

* revert since mypy won't pick the correct signature

* lint 1

* fix test

* fix test

* docs and celery config migration

* bump FAB to 4.0.0rc3, remove AUTH_STRICT_RESPONSE_CODES

* update docs for new celery config keys

* downgrade celery to 5.2.2

* ref FAB to final 4.0.0 release

* remove conflict left over

UPDATING.md

@@ -29,6 +29,9 @@ assists people when migrating to a new version.
 
 ### Breaking Changes
 
+- [19168](https://github.com/apache/superset/pull/19168): Celery upgrade to 5.X has breaking changes on it's command line invocation.
+  Please follow: https://docs.celeryq.dev/en/stable/whatsnew-5.2.html#step-1-adjust-your-command-line-invocation
+  Consider migrating you celery config if you haven't already: https://docs.celeryq.dev/en/stable/userguide/configuration.html#conf-old-settings-map
 - [19049](https://github.com/apache/superset/pull/19049): APP_ICON_WIDTH has been removed from the config. Superset should now be able to handle different logo sizes without having to explicitly set an APP_ICON_WIDTH. This might affect the size of existing custom logos as the UI will now resize them according to the specified space of maximum 148px and not according to the value of APP_ICON_WIDTH.
 - [19274](https://github.com/apache/superset/pull/19274): The `PUBLIC_ROLE_LIKE_GAMMA` config key has been removed, set `PUBLIC_ROLE_LIKE` = "Gamma" to have the same functionality.
 - [19273](https://github.com/apache/superset/pull/19273): The `SUPERSET_CELERY_WORKERS` and `SUPERSET_WORKERS` config keys has been removed. Configure celery directly using `CELERY_CONFIG` on Superset

docs/docs/installation/alerts-reports.mdx

@@ -89,12 +89,12 @@ REDIS_HOST = "redis-superset"
 REDIS_PORT = "6379"
 
 class CeleryConfig:
-    BROKER_URL = 'redis://%s:%s/0' % (REDIS_HOST, REDIS_PORT)
+    broker_url = 'redis://%s:%s/0' % (REDIS_HOST, REDIS_PORT)
-    CELERY_IMPORTS = ('superset.sql_lab', "superset.tasks", "superset.tasks.thumbnails", )
+    imports = ('superset.sql_lab', "superset.tasks", "superset.tasks.thumbnails", )
-    CELERY_RESULT_BACKEND = 'redis://%s:%s/0' % (REDIS_HOST, REDIS_PORT)
+    result_backend = 'redis://%s:%s/0' % (REDIS_HOST, REDIS_PORT)
-    CELERYD_PREFETCH_MULTIPLIER = 10
+    worker_prefetch_multiplier = 10
-    CELERY_ACKS_LATE = True
+    task_acks_late = True
-    CELERY_ANNOTATIONS = {
+    task_annotations = {
         'sql_lab.get_sql_results': {
             'rate_limit': '100/s',
         },
@@ -105,7 +105,7 @@ class CeleryConfig:
             'ignore_result': True,
         },
     }
-    CELERYBEAT_SCHEDULE = {
+    beat_schedule = {
         'reports.scheduler': {
             'task': 'reports.scheduler',
             'schedule': crontab(minute='*', hour='*'),

docs/docs/installation/async-queries-celery.mdx

@@ -23,16 +23,16 @@ and web server processes should have the same configuration.
 
 ```python
 class CeleryConfig(object):
-    BROKER_URL = 'redis://localhost:6379/0'
+    broker_url = 'redis://localhost:6379/0'
-    CELERY_IMPORTS = (
+    imports = (
         'superset.sql_lab',
         'superset.tasks',
     )
-    CELERY_RESULT_BACKEND = 'redis://localhost:6379/0'
+    result_backend = 'redis://localhost:6379/0'
-    CELERYD_LOG_LEVEL = 'DEBUG'
+    worker_log_level = 'DEBUG'
-    CELERYD_PREFETCH_MULTIPLIER = 10
+    worker_prefetch_multiplier = 10
-    CELERY_ACKS_LATE = True
+    task_acks_late = True
-    CELERY_ANNOTATIONS = {
+    task_annotations = {
         'sql_lab.get_sql_results': {
             'rate_limit': '100/s',
         },
@@ -43,7 +43,7 @@ class CeleryConfig(object):
             'ignore_result': True,
         },
     }
-    CELERYBEAT_SCHEDULE = {
+    beat_schedule = {
         'email_reports.schedule_hourly': {
             'task': 'email_reports.schedule_hourly',
             'schedule': crontab(minute=1, hour='*'),

docs/docs/installation/cache.mdx

@@ -85,11 +85,11 @@ from s3cache.s3cache import S3Cache
 ...
 
 class CeleryConfig(object):
-    BROKER_URL = "redis://localhost:6379/0"
+    broker_url = "redis://localhost:6379/0"
-    CELERY_IMPORTS = ("superset.sql_lab", "superset.tasks", "superset.tasks.thumbnails")
+    imports = ("superset.sql_lab", "superset.tasks", "superset.tasks.thumbnails")
-    CELERY_RESULT_BACKEND = "redis://localhost:6379/0"
+    result_backend = "redis://localhost:6379/0"
-    CELERYD_PREFETCH_MULTIPLIER = 10
+    worker_prefetch_multiplier = 10
-    CELERY_ACKS_LATE = True
+    task_acks_late = True
 
 
 CELERY_CONFIG = CeleryConfig

docs/docs/installation/running-on-kubernetes.mdx

@@ -344,12 +344,10 @@ configOverrides:
     from celery.schedules import crontab
 
     class CeleryConfig(object):
-      BROKER_URL = f"redis://{env('REDIS_HOST')}:{env('REDIS_PORT')}/0"
+      broker_url = f"redis://{env('REDIS_HOST')}:{env('REDIS_PORT')}/0"
-      CELERY_IMPORTS = ('superset.sql_lab', )
+      imports = ('superset.sql_lab', "superset.tasks", "superset.tasks.thumbnails", )
-      CELERY_RESULT_BACKEND = f"redis://{env('REDIS_HOST')}:{env('REDIS_PORT')}/0"
+      result_backend = f"redis://{env('REDIS_HOST')}:{env('REDIS_PORT')}/0"
-      CELERY_ANNOTATIONS = {'tasks.add': {'rate_limit': '10/s'}}
+      task_annotations = {
-      CELERY_IMPORTS = ('superset.sql_lab', "superset.tasks", "superset.tasks.thumbnails", )
-      CELERY_ANNOTATIONS = {
           'sql_lab.get_sql_results': {
               'rate_limit': '100/s',
           },
@@ -360,7 +358,7 @@ configOverrides:
               'ignore_result': True,
           },
       }
-      CELERYBEAT_SCHEDULE = {
+      beat_schedule = {
           'reports.scheduler': {
               'task': 'reports.scheduler',
               'schedule': crontab(minute='*', hour='*'),

requirements/base.txt

@@ -11,7 +11,7 @@ aiohttp==3.7.4.post0
     # via slackclient
 alembic==1.6.5
     # via flask-migrate
-amqp==2.6.1
+amqp==5.1.0
     # via kombu
 apispec[yaml]==3.3.2
     # via flask-appbuilder
@@ -33,17 +33,27 @@ brotli==1.0.9
     # via flask-compress
 cachelib==0.4.1
     # via apache-superset
-celery==4.4.7
+celery==5.2.2
     # via apache-superset
 cffi==1.14.6
     # via cryptography
 chardet==4.0.0
     # via aiohttp
-click==7.1.2
+click==8.0.4
     # via
     #   apache-superset
+    #   celery
+    #   click-didyoumean
+    #   click-plugins
+    #   click-repl
     #   flask
     #   flask-appbuilder
+click-didyoumean==0.3.0
+    # via celery
+click-plugins==1.1.1
+    # via celery
+click-repl==0.2.0
+    # via celery
 colorama==0.4.4
     # via
     #   apache-superset
@@ -56,15 +66,13 @@ croniter==1.0.15
     # via apache-superset
 cryptography==3.4.7
     # via apache-superset
-defusedxml==0.7.1
-    # via python3-openid
 deprecation==2.1.0
     # via apache-superset
 dnspython==2.1.0
     # via email-validator
 email-validator==1.1.3
     # via flask-appbuilder
-flask==1.1.4
+flask==2.0.3
     # via
     #   apache-superset
     #   flask-appbuilder
@@ -74,10 +82,9 @@ flask==1.1.4
     #   flask-jwt-extended
     #   flask-login
     #   flask-migrate
-    #   flask-openid
     #   flask-sqlalchemy
     #   flask-wtf
-flask-appbuilder==3.4.5
+flask-appbuilder==4.0.0
     # via apache-superset
 flask-babel==1.0.0
     # via flask-appbuilder
@@ -85,14 +92,12 @@ flask-caching==1.10.1
     # via apache-superset
 flask-compress==1.10.1
     # via apache-superset
-flask-jwt-extended==3.25.1
+flask-jwt-extended==4.3.1
     # via flask-appbuilder
 flask-login==0.4.1
     # via flask-appbuilder
 flask-migrate==3.1.0
     # via apache-superset
-flask-openid==1.3.0
-    # via flask-appbuilder
 flask-sqlalchemy==2.5.1
     # via
     #   flask-appbuilder
@@ -123,18 +128,17 @@ idna==3.2
     #   yarl
 isodate==0.6.0
     # via apache-superset
-itsdangerous==1.1.0
+itsdangerous==2.1.1
     # via
-    #   apache-superset
     #   flask
     #   flask-wtf
-jinja2==2.11.3
+jinja2==3.0.3
     # via
     #   flask
     #   flask-babel
 jsonschema==3.2.0
     # via flask-appbuilder
-kombu==4.6.11
+kombu==5.2.4
     # via celery
 korean-lunar-calendar==0.2.1
     # via holidays
@@ -180,11 +184,13 @@ polyline==1.4.0
     # via apache-superset
 prison==0.2.1
     # via flask-appbuilder
+prompt-toolkit==3.0.28
+    # via click-repl
 pyarrow==5.0.0
     # via apache-superset
 pycparser==2.20
     # via cffi
-pyjwt==1.7.1
+pyjwt==2.2.0
     # via
     #   apache-superset
     #   flask-appbuilder
@@ -213,9 +219,7 @@ python-editor==1.0.4
     # via alembic
 python-geohash==0.8.5
     # via apache-superset
-python3-openid==3.2.0
+pytz==2021.3
-    # via flask-openid
-pytz==2021.1
     # via
     #   babel
     #   celery
@@ -237,7 +241,7 @@ simplejson==3.17.3
 six==1.16.0
     # via
     #   bleach
-    #   flask-jwt-extended
+    #   click-repl
     #   flask-talisman
     #   holidays
     #   isodate
@@ -273,13 +277,16 @@ typing-extensions==3.10.0.0
     #   apache-superset
 urllib3==1.26.6
     # via selenium
-vine==1.3.0
+vine==5.0.0
     # via
     #   amqp
     #   celery
+    #   kombu
+wcwidth==0.2.5
+    # via prompt-toolkit
 webencodings==0.5.1
     # via bleach
-werkzeug==1.0.1
+werkzeug==2.0.3
     # via
     #   flask
     #   flask-jwt-extended

requirements/integration.in

@@ -18,4 +18,4 @@ pip-compile-multi!=1.5.9
 pre-commit
 tox
 py>=1.10.0
-click==7.1.2
+click

requirements/integration.txt

@@ -1,4 +1,4 @@
-# SHA1:03eb2d96afe21f1bda1ab33b4cf84e670a1efe21
+# SHA1:8e2dd1e795bcad7451376b3653eb03465e4f05d3
 #
 # This file is autogenerated by pip-compile-multi
 # To update, run:
@@ -9,7 +9,7 @@ backports.entry-points-selectable==1.1.0
     # via virtualenv
 cfgv==3.3.0
     # via pre-commit
-click==7.1.2
+click==8.0.4
     # via
     #   -r requirements/integration.in
     #   pip-compile-multi

requirements/testing.txt

@@ -116,8 +116,6 @@ pexpect==4.8.0
     # via ipython
 pickleshare==0.7.5
     # via ipython
-prompt-toolkit==3.0.19
-    # via ipython
 proto-plus==1.19.7
     # via
     #   google-cloud-bigquery
@@ -178,8 +176,6 @@ trino==0.306
     # via sqlalchemy-trino
 typing-inspect==0.7.1
     # via libcst
-wcwidth==0.2.5
-    # via prompt-toolkit
 websocket-client==1.2.0
     # via docker
 

scripts/tests/run.sh

@@ -138,5 +138,5 @@ fi
 
 if [ $RUN_TESTS -eq 1 ]
 then
-  pytest -x -s "${TEST_MODULE}"
+  pytest --durations=0 --maxfail=1 "${TEST_MODULE}"
 fi

setup.cfg

@@ -41,7 +41,7 @@ disallow_untyped_calls = true
 disallow_untyped_defs = true
 ignore_missing_imports = true
 no_implicit_optional = true
-warn_unused_ignores = true
+warn_unused_ignores = false
 
 [mypy-superset.migrations.versions.*]
 ignore_errors = true

setup.py

@@ -70,15 +70,15 @@ setup(
         "backoff>=1.8.0",
         "bleach>=3.0.2, <4.0.0",
         "cachelib>=0.4.1,<0.5",
-        "celery>=4.3.0, <5.0.0, !=4.4.1",
+        "celery>=5.2.2, <6.0.0",
-        "click<8",
+        "click>=8.0.3",
         "colorama",
         "croniter>=0.3.28",
         "cron-descriptor",
         "cryptography>=3.3.2",
         "deprecation>=2.1.0, <2.2.0",
-        "flask>=1.1.0, <2.0.0",
+        "flask>=2.0.0, <3.0.0",
-        "flask-appbuilder>=3.4.5, <4.0.0",
+        "flask-appbuilder>=4.0.0, <5.0.0",
         "flask-caching>=1.10.0",
         "flask-compress",
         "flask-talisman",
@@ -90,7 +90,6 @@ setup(
         "gunicorn>=20.1.0",
         "holidays==0.10.3",  # PINNED! https://github.com/dr-prodigy/python-holidays/issues/406
         "humanize",
-        "itsdangerous>=1.0.0, <2.0.0",  # https://github.com/apache/superset/pull/14627
         "isodate",
         "markdown>=3.0",
         "msgpack>=1.0.0, <1.1",
@@ -104,7 +103,7 @@ setup(
         "python-geohash",
         "pyarrow>=5.0.0, <6.0",
         "pyyaml>=5.4",
-        "PyJWT>=1.7.1, <2",
+        "PyJWT>=2.0.0, <2.3.0",
         "redis",
         "selenium>=3.141.0",
         "simplejson>=3.15.0",

superset/cli/importexport.py

@@ -64,7 +64,9 @@ if feature_flags.get("VERSIONED_EXPORT"):
         from superset.dashboards.commands.export import ExportDashboardsCommand
         from superset.models.dashboard import Dashboard
 
-        g.user = security_manager.find_user(username="admin")
+        g.user = security_manager.find_user(  # pylint: disable=assigning-non-slot
+            username="admin"
+        )
 
         dashboard_ids = [id_ for (id_,) in db.session.query(Dashboard.id).all()]
         timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
@@ -96,7 +98,9 @@ if feature_flags.get("VERSIONED_EXPORT"):
         from superset.connectors.sqla.models import SqlaTable
         from superset.datasets.commands.export import ExportDatasetsCommand
 
-        g.user = security_manager.find_user(username="admin")
+        g.user = security_manager.find_user(  # pylint: disable=assigning-non-slot
+            username="admin"
+        )
 
         dataset_ids = [id_ for (id_,) in db.session.query(SqlaTable.id).all()]
         timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
@@ -135,7 +139,9 @@ if feature_flags.get("VERSIONED_EXPORT"):
         )
 
         if username is not None:
-            g.user = security_manager.find_user(username=username)
+            g.user = security_manager.find_user(  # pylint: disable=assigning-non-slot
+                username=username
+            )
         if is_zipfile(path):
             with ZipFile(path) as bundle:
                 contents = get_contents_from_bundle(bundle)
@@ -299,7 +305,9 @@ else:
         elif path_object.exists() and recursive:
             files.extend(path_object.rglob("*.json"))
         if username is not None:
-            g.user = security_manager.find_user(username=username)
+            g.user = security_manager.find_user(  # pylint: disable=assigning-non-slot
+                username=username
+            )
         contents = {}
         for path_ in files:
             with open(path_) as file:

superset/config.py

@@ -300,8 +300,6 @@ AUTH_TYPE = AUTH_DB
 #    { 'name': 'Yahoo', 'url': 'https://open.login.yahoo.com/' },
 #    { 'name': 'Flickr', 'url': 'https://www.flickr.com/<username>' },
 
-AUTH_STRICT_RESPONSE_CODES = True
-
 # ---------------------------------------------------
 # Roles config
 # ---------------------------------------------------
@@ -747,13 +745,13 @@ DASHBOARD_AUTO_REFRESH_MODE: Literal["fetch", "force"] = "force"
 
 
 class CeleryConfig:  # pylint: disable=too-few-public-methods
-    BROKER_URL = "sqla+sqlite:///celerydb.sqlite"
+    broker_url = "sqla+sqlite:///celerydb.sqlite"
-    CELERY_IMPORTS = ("superset.sql_lab", "superset.tasks")
+    imports = ("superset.sql_lab", "superset.tasks")
-    CELERY_RESULT_BACKEND = "db+sqlite:///celery_results.sqlite"
+    result_backend = "db+sqlite:///celery_results.sqlite"
-    CELERYD_LOG_LEVEL = "DEBUG"
+    worker_log_level = "DEBUG"
-    CELERYD_PREFETCH_MULTIPLIER = 1
+    worker_prefetch_multiplier = 1
-    CELERY_ACKS_LATE = False
+    task_acks_late = False
-    CELERY_ANNOTATIONS = {
+    task_annotations = {
         "sql_lab.get_sql_results": {"rate_limit": "100/s"},
         "email_reports.send": {
             "rate_limit": "1/s",
@@ -762,7 +760,7 @@ class CeleryConfig:  # pylint: disable=too-few-public-methods
             "ignore_result": True,
         },
     }
-    CELERYBEAT_SCHEDULE = {
+    beat_schedule = {
         "email_reports.schedule_hourly": {
             "task": "email_reports.schedule_hourly",
             "schedule": crontab(minute=1, hour="*"),

superset/tasks/async_queries.py

@@ -47,13 +47,17 @@ query_timeout = current_app.config[
 def ensure_user_is_set(user_id: Optional[int]) -> None:
     user_is_not_set = not (hasattr(g, "user") and g.user is not None)
     if user_is_not_set and user_id is not None:
-        g.user = security_manager.get_user_by_id(user_id)
+        g.user = security_manager.get_user_by_id(  # pylint: disable=assigning-non-slot
+            user_id
+        )
     elif user_is_not_set:
-        g.user = security_manager.get_anonymous_user()
+        g.user = (  # pylint: disable=assigning-non-slot
+            security_manager.get_anonymous_user()
+        )
 
 
 def set_form_data(form_data: Dict[str, Any]) -> None:
-    g.form_data = form_data
+    g.form_data = form_data  # pylint: disable=assigning-non-slot
 
 
 def _create_query_context_from_form(form_data: Dict[str, Any]) -> QueryContext:

superset/utils/async_query_manager.py

@@ -134,7 +134,11 @@ class AsyncQueryManager:
                 session["async_user_id"] = user_id
 
                 sub = str(user_id) if user_id else None
-                token = self.generate_jwt({"channel": async_channel_id, "sub": sub})
+                token = jwt.encode(
+                    {"channel": async_channel_id, "sub": sub},
+                    self._jwt_secret,
+                    algorithm="HS256",
+                )
 
                 response.set_cookie(
                     self._jwt_cookie_name,
@@ -146,21 +150,13 @@ class AsyncQueryManager:
 
             return response
 
-    def generate_jwt(self, data: Dict[str, Any]) -> str:
-        encoded_jwt = jwt.encode(data, self._jwt_secret, algorithm="HS256")
-        return encoded_jwt.decode("utf-8")
-
-    def parse_jwt(self, token: str) -> Dict[str, Any]:
-        data = jwt.decode(token, self._jwt_secret, algorithms=["HS256"])
-        return data
-
     def parse_jwt_from_request(self, req: Request) -> Dict[str, Any]:
         token = req.cookies.get(self._jwt_cookie_name)
         if not token:
             raise AsyncQueryTokenException("Token not preset")
 
         try:
-            return self.parse_jwt(token)
+            return jwt.decode(token, self._jwt_secret, algorithms=["HS256"])
         except Exception as ex:
             logger.warning(ex)
             raise AsyncQueryTokenException("Failed to parse token") from ex

superset/views/core.py

@@ -1841,7 +1841,7 @@ class Superset(BaseSupersetView):  # pylint: disable=too-many-public-methods
                     force=True,
                 )
 
-                g.form_data = form_data
+                g.form_data = form_data  # pylint: disable=assigning-non-slot
                 payload = obj.get_payload()
                 delattr(g, "form_data")
                 error = payload["errors"] or None

tests/integration_tests/core_tests.py

@@ -1254,8 +1254,10 @@ class TestCore(SupersetTestCase):
         self.assertEqual(resp.status_code, 403)
 
     @mock.patch("superset.views.core.results_backend_use_msgpack", False)
-    @mock.patch("superset.views.core.results_backend")
+    def test_display_limit(self):
-    def test_display_limit(self, mock_results_backend):
+        from superset.views import core
+
+        core.results_backend = mock.Mock()
         self.login()
 
         data = [{"col_0": i} for i in range(100)]
@@ -1284,7 +1286,7 @@ class TestCore(SupersetTestCase):
         app.config["RESULTS_BACKEND_USE_MSGPACK"] = False
         serialized_payload = sql_lab._serialize_payload(payload, False)
         compressed = utils.zlib_compress(serialized_payload)
-        mock_results_backend.get.return_value = compressed
+        core.results_backend.get.return_value = compressed
 
         with mock.patch("superset.views.core.db") as mock_superset_db:
             mock_superset_db.session.query().filter_by().one_or_none.return_value = (

tests/integration_tests/dashboards/api_tests.py

@@ -1348,10 +1348,8 @@ class TestDashboardApi(SupersetTestCase, ApiOwnersTestCaseMixin, InsertChartMixi
         dashboards_ids = get_dashboards_ids(db, ["world_health", "births"])
         uri = f"api/v1/dashboard/export/?q={prison.dumps(dashboards_ids)}"
 
-        # freeze time to ensure filename is deterministic
+        rv = self.get_assert_metric(uri, "export")
-        with freeze_time("2020-01-01T00:00:00Z"):
+        headers = generate_download_headers("json")["Content-Disposition"]
-            rv = self.get_assert_metric(uri, "export")
-            headers = generate_download_headers("json")["Content-Disposition"]
 
         assert rv.status_code == 200
         assert rv.headers["Content-Disposition"] == headers

tests/integration_tests/security/api_tests.py

@@ -92,7 +92,10 @@ class TestSecurityGuestTokenApi(SupersetTestCase):
         self.assert200(response)
         token = json.loads(response.data)["token"]
         decoded_token = jwt.decode(
-            token, self.app.config["GUEST_TOKEN_JWT_SECRET"], audience=get_url_host()
+            token,
+            self.app.config["GUEST_TOKEN_JWT_SECRET"],
+            audience=get_url_host(),
+            algorithms=["HS256"],
         )
         self.assertEqual(user, decoded_token["user"])
         self.assertEqual(resource, decoded_token["resources"][0])

tests/integration_tests/sqllab_tests.py

@@ -24,9 +24,9 @@ from celery.exceptions import SoftTimeLimitExceeded
 from parameterized import parameterized
 from random import random
 from unittest import mock
-from superset.extensions import db
 import prison
 
+from freezegun import freeze_time
 from superset import db, security_manager
 from superset.connectors.sqla.models import SqlaTable
 from superset.db_engine_specs import BaseEngineSpec
@@ -34,16 +34,12 @@ from superset.db_engine_specs.hive import HiveEngineSpec
 from superset.db_engine_specs.presto import PrestoEngineSpec
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import SupersetErrorException
-from superset.models.core import Database
 from superset.models.sql_lab import Query, SavedQuery
 from superset.result_set import SupersetResultSet
 from superset.sqllab.limiting_factor import LimitingFactor
 from superset.sql_lab import (
     cancel_query,
     execute_sql_statements,
-    execute_sql_statement,
-    get_sql_results,
-    SqlLabException,
     apply_limit_if_exists,
 )
 from superset.sql_parse import CtasMethod
@@ -157,8 +153,6 @@ class TestSqlLab(SupersetTestCase):
         """
         SQLLab: Test SQLLab query execution info propagation to saved queries
         """
-        from freezegun import freeze_time
-
         self.login("admin")
 
         sql_statement = "SELECT * FROM birth_names LIMIT 10"
@@ -167,7 +161,7 @@ class TestSqlLab(SupersetTestCase):
         db.session.add(saved_query)
         db.session.commit()
 
-        with freeze_time("2020-01-01T00:00:00Z"):
+        with freeze_time(datetime.now().isoformat(timespec="seconds")):
             self.run_sql(sql_statement, "1")
             saved_query_ = (
                 db.session.query(SavedQuery)
@@ -178,9 +172,9 @@ class TestSqlLab(SupersetTestCase):
             )
             assert saved_query_.rows is not None
             assert saved_query_.last_run == datetime.now()
-            # Rollback changes
+        # Rollback changes
-            db.session.delete(saved_query_)
+        db.session.delete(saved_query_)
-            db.session.commit()
+        db.session.commit()
 
     @parameterized.expand([CtasMethod.TABLE, CtasMethod.VIEW])
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")