Skip to content

Commit

Permalink
chore: add statsd support to base API and refactor (#22887)
Browse files Browse the repository at this point in the history
  • Loading branch information
dpgaspar committed Jan 27, 2023
1 parent bed10a0 commit d00ba15
Show file tree
Hide file tree
Showing 18 changed files with 167 additions and 138 deletions.
6 changes: 3 additions & 3 deletions superset/advanced_data_type/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand All @@ -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"

Expand Down
8 changes: 3 additions & 5 deletions superset/async_events/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 5 additions & 4 deletions superset/available_domains/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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,
Expand Down
6 changes: 4 additions & 2 deletions superset/cachekeys/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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),
Expand Down
14 changes: 4 additions & 10 deletions superset/dashboards/permalink/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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"
Expand Down
5 changes: 3 additions & 2 deletions superset/databases/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand Down
9 changes: 5 additions & 4 deletions superset/explore/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
Expand All @@ -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,
Expand Down
18 changes: 8 additions & 10 deletions superset/explore/form_data/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
16 changes: 6 additions & 10 deletions superset/explore/permalink/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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"
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 3 additions & 2 deletions superset/extensions/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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()
31 changes: 31 additions & 0 deletions superset/extensions/stats_logger.py
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit d00ba15

Please sign in to comment.