chore: add statsd support to base API and refactor (#22887)

superset/advanced_data_type/api.py

@@ -18,7 +18,7 @@
 
 from flask import current_app as app
 from flask.wrappers import Response
-from flask_appbuilder.api import BaseApi, expose, permission_name, protect, rison, safe
+from flask_appbuilder.api import expose, permission_name, protect, rison, safe
 from flask_babel import lazy_gettext as _
 
 from superset.advanced_data_type.schemas import (
@@ -27,12 +27,13 @@
 )
 from superset.advanced_data_type.types import AdvancedDataTypeResponse
 from superset.extensions import event_logger
+from superset.views.base_api import BaseSupersetApi
 
 config = app.config
 ADVANCED_DATA_TYPES = config["ADVANCED_DATA_TYPES"]
 
 
-class AdvancedDataTypeRestApi(BaseApi):
+class AdvancedDataTypeRestApi(BaseSupersetApi):
     """
     Advanced Data Type Rest API
     -Will return available AdvancedDataTypes when the /types endpoint is accessed
@@ -41,7 +42,6 @@ class AdvancedDataTypeRestApi(BaseApi):
     """
 
     allow_browser_login = True
-    include_route_methods = {"get", "get_types"}
     resource_name = "advanced_data_type"
     class_permission_name = "AdvancedDataType"
 

superset/async_events/api.py

@@ -18,21 +18,19 @@
 
 from flask import request, Response
 from flask_appbuilder import expose
-from flask_appbuilder.api import BaseApi, safe
+from flask_appbuilder.api import safe
 from flask_appbuilder.security.decorators import permission_name, protect
 
 from superset.extensions import async_query_manager, event_logger
 from superset.utils.async_query_manager import AsyncQueryTokenException
+from superset.views.base_api import BaseSupersetApi
 
 logger = logging.getLogger(__name__)
 
 
-class AsyncEventsRestApi(BaseApi):
+class AsyncEventsRestApi(BaseSupersetApi):
     resource_name = "async_event"
     allow_browser_login = True
-    include_route_methods = {
-        "events",
-    }
 
     @expose("/", methods=["GET"])
     @event_logger.log_this

superset/available_domains/api.py

@@ -17,21 +17,21 @@
 import logging
 
 from flask import Response
-from flask_appbuilder.api import BaseApi, expose, protect, safe
+from flask_appbuilder.api import expose, protect, safe
 
 from superset import conf
 from superset.available_domains.schemas import AvailableDomainsSchema
-from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
+from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP
 from superset.extensions import event_logger
+from superset.views.base_api import BaseSupersetApi, statsd_metrics
 
 logger = logging.getLogger(__name__)
 
 
-class AvailableDomainsRestApi(BaseApi):
+class AvailableDomainsRestApi(BaseSupersetApi):
     available_domains_schema = AvailableDomainsSchema()
 
     method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP
-    include_route_methods = {RouteMethod.GET}
     allow_browser_login = True
     class_permission_name = "AvailableDomains"
     resource_name = "available_domains"
@@ -41,6 +41,7 @@ class AvailableDomainsRestApi(BaseApi):
     @expose("/", methods=["GET"])
     @protect()
     @safe
+    @statsd_metrics
     @event_logger.log_this_with_context(
         action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.get",
         log_to_statsd=True,

superset/cachekeys/api.py

@@ -26,7 +26,7 @@
 
 from superset.cachekeys.schemas import CacheInvalidationRequestSchema
 from superset.connectors.sqla.models import SqlaTable
-from superset.extensions import cache_manager, db, event_logger
+from superset.extensions import cache_manager, db, event_logger, stats_logger_manager
 from superset.models.cache import CacheKey
 from superset.views.base_api import BaseSupersetModelRestApi, statsd_metrics
 
@@ -117,7 +117,9 @@ def invalidate(self) -> Response:
                 )
                 db.session.execute(delete_stmt)
                 db.session.commit()
-                self.stats_logger.gauge("invalidated_cache", len(cache_keys))
+                stats_logger_manager.instance.gauge(
+                    "invalidated_cache", len(cache_keys)
+                )
                 logger.info(
                     "Invalidated %s cache records for %s datasources",
                     len(cache_keys),

superset/dashboards/permalink/api.py

@@ -17,10 +17,10 @@
 import logging
 
 from flask import request, Response
-from flask_appbuilder.api import BaseApi, expose, protect, safe
+from flask_appbuilder.api import expose, protect, safe
 from marshmallow import ValidationError
 
-from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
+from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP
 from superset.dashboards.commands.exceptions import (
     DashboardAccessDeniedError,
     DashboardNotFoundError,
@@ -33,20 +33,14 @@
 from superset.dashboards.permalink.schemas import DashboardPermalinkPostSchema
 from superset.extensions import event_logger
 from superset.key_value.exceptions import KeyValueAccessDeniedError
-from superset.views.base_api import requires_json
+from superset.views.base_api import BaseSupersetApi, requires_json
 
 logger = logging.getLogger(__name__)
 
 
-class DashboardPermalinkRestApi(BaseApi):
+class DashboardPermalinkRestApi(BaseSupersetApi):
     add_model_schema = DashboardPermalinkPostSchema()
     method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP
-    include_route_methods = {
-        RouteMethod.POST,
-        RouteMethod.PUT,
-        RouteMethod.GET,
-        RouteMethod.DELETE,
-    }
     allow_browser_login = True
     class_permission_name = "DashboardPermalinkRestApi"
     resource_name = "dashboard"

superset/databases/decorators.py

@@ -21,6 +21,7 @@
 from flask import g
 from flask_babel import lazy_gettext as _
 
+from superset.extensions import stats_logger_manager
 from superset.models.core import Database
 from superset.sql_parse import Table
 from superset.utils.core import parse_js_uri_path_item
@@ -46,14 +47,14 @@ def wraps(
             return self.response_422(message=_("Table name undefined"))
         database: Database = self.datamodel.get(pk)
         if not database:
-            self.stats_logger.incr(
+            stats_logger_manager.instance.incr(
                 f"database_not_found_{self.__class__.__name__}.select_star"
             )
             return self.response_404()
         if not self.appbuilder.sm.can_access_table(
             database, Table(table_name_parsed, schema_name_parsed)
         ):
-            self.stats_logger.incr(
+            stats_logger_manager.instance.incr(
                 f"permisssion_denied_{self.__class__.__name__}.select_star"
             )
             logger.warning(

superset/explore/api.py

@@ -17,10 +17,10 @@
 import logging
 
 from flask import g, request, Response
-from flask_appbuilder.api import BaseApi, expose, protect, safe
+from flask_appbuilder.api import expose, protect, safe
 
 from superset.charts.commands.exceptions import ChartNotFoundError
-from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
+from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP
 from superset.explore.commands.get import GetExploreCommand
 from superset.explore.commands.parameters import CommandParameters
 from superset.explore.exceptions import DatasetAccessDeniedError, WrongEndpointError
@@ -31,13 +31,13 @@
     TemporaryCacheAccessDeniedError,
     TemporaryCacheResourceNotFoundError,
 )
+from superset.views.base_api import BaseSupersetApi, statsd_metrics
 
 logger = logging.getLogger(__name__)
 
 
-class ExploreRestApi(BaseApi):
+class ExploreRestApi(BaseSupersetApi):
     method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP
-    include_route_methods = {RouteMethod.GET}
     allow_browser_login = True
     class_permission_name = "Explore"
     resource_name = "explore"
@@ -47,6 +47,7 @@ class ExploreRestApi(BaseApi):
     @expose("/", methods=["GET"])
     @protect()
     @safe
+    @statsd_metrics
     @event_logger.log_this_with_context(
         action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.get",
         log_to_statsd=True,

superset/explore/form_data/api.py

@@ -17,10 +17,10 @@
 import logging
 
 from flask import request, Response
-from flask_appbuilder.api import BaseApi, expose, protect, safe
+from flask_appbuilder.api import expose, protect, safe
 from marshmallow import ValidationError
 
-from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
+from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP
 from superset.explore.form_data.commands.create import CreateFormDataCommand
 from superset.explore.form_data.commands.delete import DeleteFormDataCommand
 from superset.explore.form_data.commands.get import GetFormDataCommand
@@ -32,21 +32,15 @@
     TemporaryCacheAccessDeniedError,
     TemporaryCacheResourceNotFoundError,
 )
-from superset.views.base_api import requires_json
+from superset.views.base_api import BaseSupersetApi, requires_json, statsd_metrics
 
 logger = logging.getLogger(__name__)
 
 
-class ExploreFormDataRestApi(BaseApi):
+class ExploreFormDataRestApi(BaseSupersetApi):
     add_model_schema = FormDataPostSchema()
     edit_model_schema = FormDataPutSchema()
     method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP
-    include_route_methods = {
-        RouteMethod.POST,
-        RouteMethod.PUT,
-        RouteMethod.GET,
-        RouteMethod.DELETE,
-    }
     allow_browser_login = True
     class_permission_name = "ExploreFormDataRestApi"
     resource_name = "explore"
@@ -56,6 +50,7 @@ class ExploreFormDataRestApi(BaseApi):
     @expose("/form_data", methods=["POST"])
     @protect()
     @safe
+    @statsd_metrics
     @event_logger.log_this_with_context(
         action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.post",
         log_to_statsd=False,
@@ -120,6 +115,7 @@ def post(self) -> Response:
     @expose("/form_data/<string:key>", methods=["PUT"])
     @protect()
     @safe
+    @statsd_metrics
     @event_logger.log_this_with_context(
         action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.put",
         log_to_statsd=True,
@@ -193,6 +189,7 @@ def put(self, key: str) -> Response:
     @expose("/form_data/<string:key>", methods=["GET"])
     @protect()
     @safe
+    @statsd_metrics
     @event_logger.log_this_with_context(
         action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.get",
         log_to_statsd=True,
@@ -244,6 +241,7 @@ def get(self, key: str) -> Response:
     @expose("/form_data/<string:key>", methods=["DELETE"])
     @protect()
     @safe
+    @statsd_metrics
     @event_logger.log_this_with_context(
         action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.delete",
         log_to_statsd=True,

superset/explore/permalink/api.py

@@ -17,14 +17,14 @@
 import logging
 
 from flask import request, Response
-from flask_appbuilder.api import BaseApi, expose, protect, safe
+from flask_appbuilder.api import expose, protect, safe
 from marshmallow import ValidationError
 
 from superset.charts.commands.exceptions import (
     ChartAccessDeniedError,
     ChartNotFoundError,
 )
-from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
+from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP
 from superset.datasets.commands.exceptions import (
     DatasetAccessDeniedError,
     DatasetNotFoundError,
@@ -35,20 +35,14 @@
 from superset.explore.permalink.schemas import ExplorePermalinkPostSchema
 from superset.extensions import event_logger
 from superset.key_value.exceptions import KeyValueAccessDeniedError
-from superset.views.base_api import requires_json
+from superset.views.base_api import BaseSupersetApi, requires_json, statsd_metrics
 
 logger = logging.getLogger(__name__)
 
 
-class ExplorePermalinkRestApi(BaseApi):
+class ExplorePermalinkRestApi(BaseSupersetApi):
     add_model_schema = ExplorePermalinkPostSchema()
     method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP
-    include_route_methods = {
-        RouteMethod.POST,
-        RouteMethod.PUT,
-        RouteMethod.GET,
-        RouteMethod.DELETE,
-    }
     allow_browser_login = True
     class_permission_name = "ExplorePermalinkRestApi"
     resource_name = "explore"
@@ -58,6 +52,7 @@ class ExplorePermalinkRestApi(BaseApi):
     @expose("/permalink", methods=["POST"])
     @protect()
     @safe
+    @statsd_metrics
     @event_logger.log_this_with_context(
         action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.post",
         log_to_statsd=False,
@@ -118,6 +113,7 @@ def post(self) -> Response:
     @expose("/permalink/<string:key>", methods=["GET"])
     @protect()
     @safe
+    @statsd_metrics
     @event_logger.log_this_with_context(
         action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.get",
         log_to_statsd=False,

superset/extensions/__init__.py

@@ -16,7 +16,6 @@
 # under the License.
 import json
 import os
-from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional
 
 import celery
@@ -29,6 +28,7 @@
 from werkzeug.local import LocalProxy
 
 from superset.extensions.ssh import SSHManagerFactory
+from superset.extensions.stats_logger import BaseStatsLoggerManager
 from superset.utils.async_query_manager import AsyncQueryManager
 from superset.utils.cache_manager import CacheManager
 from superset.utils.encrypt import EncryptedFieldFactory
@@ -127,5 +127,6 @@ def init_app(self, app: Flask) -> None:
 profiling = ProfilingExtension()
 results_backend_manager = ResultsBackendManager()
 security_manager = LocalProxy(lambda: appbuilder.sm)
-talisman = Talisman()
 ssh_manager_factory = SSHManagerFactory()
+stats_logger_manager = BaseStatsLoggerManager()
+talisman = Talisman()

superset/extensions/stats_logger.py

@@ -0,0 +1,31 @@
+# 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 flask import Flask
+
+from superset.stats_logger import BaseStatsLogger
+
+
+class BaseStatsLoggerManager:
+    def __init__(self) -> None:
+        self._stats_logger = BaseStatsLogger()
+
+    def init_app(self, app: Flask) -> None:
+        self._stats_logger = app.config["STATS_LOGGER"]
+
+    @property
+    def instance(self) -> BaseStatsLogger:
+        return self._stats_logger