feat: add profiling to Superset pages

requirements/development.in

@@ -26,3 +26,4 @@ tableschema
 thrift>=0.11.0,<1.0.0
 pygithub>=1.54.1,<2.0.0
 progress>=1.5,<2
+pyinstrument>=4.0.2,<5

requirements/development.txt

@@ -1,4 +1,4 @@
-# SHA1:c470411e2e9cb04b412a94f80a6a9d870bece74d
+# SHA1:1144991012e228fb2ef85afbf78a635e7d5a33f1
 #
 # This file is autogenerated by pip-compile-multi
 # To update, run:
@@ -54,6 +54,8 @@ pygithub==1.54.1
     # via -r requirements/development.in
 pyhive[hive]==0.6.3
     # via -r requirements/development.in
+pyinstrument==4.0.2
+    # via -r requirements/development.in
 requests==2.24.0
     # via
     #   pydruid

setup.cfg

@@ -30,7 +30,7 @@ combine_as_imports = true
 include_trailing_comma = true
 line_length = 88
 known_first_party = superset
-known_third_party =alembic,apispec,backoff,bleach,cachelib,celery,click,colorama,cron_descriptor,croniter,cryptography,dateutil,deprecation,flask,flask_appbuilder,flask_babel,flask_caching,flask_compress,flask_jwt_extended,flask_login,flask_migrate,flask_sqlalchemy,flask_talisman,flask_testing,flask_wtf,freezegun,geohash,geopy,graphlib,holidays,humanize,isodate,jinja2,jwt,markdown,markupsafe,marshmallow,marshmallow_enum,msgpack,numpy,pandas,parameterized,parsedatetime,pgsanity,pkg_resources,polyline,prison,progress,pyarrow,pyhive,pyparsing,pytest,pytest_mock,pytz,redis,requests,selenium,setuptools,simplejson,slack,sqlalchemy,sqlalchemy_utils,sqlparse,tabulate,typing_extensions,werkzeug,wtforms,wtforms_json,yaml
+known_third_party =alembic,apispec,backoff,bleach,cachelib,celery,click,colorama,cron_descriptor,croniter,cryptography,dateutil,deprecation,flask,flask_appbuilder,flask_babel,flask_caching,flask_compress,flask_jwt_extended,flask_login,flask_migrate,flask_sqlalchemy,flask_talisman,flask_testing,flask_wtf,freezegun,geohash,geopy,graphlib,holidays,humanize,isodate,jinja2,jwt,markdown,markupsafe,marshmallow,marshmallow_enum,msgpack,numpy,pandas,parameterized,parsedatetime,pgsanity,pkg_resources,polyline,prison,progress,pyarrow,pyhive,pyinstrument,pyparsing,pytest,pytest_mock,pytz,redis,requests,selenium,setuptools,simplejson,slack,sqlalchemy,sqlalchemy_utils,sqlparse,tabulate,typing_extensions,werkzeug,wtforms,wtforms_json,yaml
 multi_line_output = 3
 order_by_type = false
 

superset/app.py

@@ -21,6 +21,7 @@
 from flask import Flask
 
 from superset.initialization import SupersetAppInitializer
+from superset.utils.profiler import SupersetProfiler
 
 logger = logging.getLogger(__name__)
 
@@ -36,6 +37,10 @@ def create_app() -> Flask:
         app_initializer = app.config.get("APP_INITIALIZER", SupersetAppInitializer)(app)
         app_initializer.init_app()
 
+        profiling = app.config.get("PROFILING", False)
+        if profiling:
+            app.wsgi_app = SupersetProfiler(app.wsgi_app)  # type: ignore
+
         return app
 
     # Make sure that bootstrap errors ALWAYS get logged

superset/config.py

@@ -196,6 +196,10 @@ def _try_json_readsha(  # pylint: disable=unused-argument
 DEBUG = os.environ.get("FLASK_ENV") == "development"
 FLASK_USE_RELOAD = True
 
+# Enable profiling of Python calls. Turn this on and append ``?_instrument=1``
+# to the page to see the call stack.
+PROFILING = False
+
 # Superset allows server-side python stacktraces to be surfaced to the
 # user when this feature is on. This may has security implications
 # and it's more secure to turn it off in production settings.

superset/utils/profiler.py

@@ -0,0 +1,52 @@
+# 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.
+
+from typing import Any, Callable
+from unittest import mock
+
+from pyinstrument import Profiler
+from werkzeug.wrappers import Request, Response
+
+
+class SupersetProfiler:
+    """
+    WSGI middleware to instrument Superset.
+
+    To see the instrumentation for a given page, set `PROFILING=True`
+    in the config, and append `?_instrument=1` to the page.
+    """
+
+    def __init__(
+        self, app: Callable[[Any, Any], Any], interval: float = 0.0001,
+    ):
+        self.app = app
+        self.interval = interval
+
+    @Request.application
+    def __call__(self, request: Request) -> Response:
+        if request.args.get("_instrument") != "1":
+            return Response.from_app(self.app, request.environ)
+
+        profiler = Profiler(interval=self.interval)
+
+        # call original request
+        fake_start_response = mock.MagicMock()
+        with profiler:
+            self.app(request.environ, fake_start_response)
+
+        # return HTML profiling information
+        return Response(profiler.output_html(), mimetype="text/html")