Skip to content

Commit

Permalink
fix: memoized decorator memory leak (#23139)
Browse files Browse the repository at this point in the history
(cherry picked from commit 79274eb)
  • Loading branch information
dpgaspar authored and Lily Kuang committed Mar 8, 2023
1 parent 2a61e15 commit ef648ae
Show file tree
Hide file tree
Showing 12 changed files with 24 additions and 203 deletions.
7 changes: 4 additions & 3 deletions superset/connectors/sqla/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from __future__ import annotations

import logging
from functools import lru_cache
from typing import (
Any,
Callable,
Expand All @@ -40,6 +41,7 @@
from sqlalchemy.orm.exc import ObjectDeletedError
from sqlalchemy.sql.type_api import TypeEngine

from superset.constants import LRU_CACHE_MAX_SIZE
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import (
SupersetGenericDBErrorException,
Expand All @@ -49,7 +51,6 @@
from superset.result_set import SupersetResultSet
from superset.sql_parse import has_table_query, insert_rls, ParsedQuery
from superset.superset_typing import ResultSetColumnType
from superset.utils.memoized import memoized

if TYPE_CHECKING:
from superset.connectors.sqla.models import SqlaTable
Expand Down Expand Up @@ -200,12 +201,12 @@ def validate_adhoc_subquery(
return ";\n".join(str(statement) for statement in statements)


@memoized
@lru_cache(maxsize=LRU_CACHE_MAX_SIZE)
def get_dialect_name(drivername: str) -> str:
return SqlaURL.create(drivername).get_dialect().name


@memoized
@lru_cache(maxsize=LRU_CACHE_MAX_SIZE)
def get_identifier_quoter(drivername: str) -> Dict[str, Callable[[str], str]]:
return SqlaURL.create(drivername).get_dialect()().identifier_preparer.quote

Expand Down
2 changes: 2 additions & 0 deletions superset/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@
QUERY_CANCEL_KEY = "cancel_query"
QUERY_EARLY_CANCEL_KEY = "early_cancel_query"

LRU_CACHE_MAX_SIZE = 256


class RouteMethod: # pylint: disable=too-few-public-methods
"""
Expand Down
8 changes: 4 additions & 4 deletions superset/jinja_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
"""Defines the templating context for SQL Lab"""
import json
import re
from functools import partial
from functools import lru_cache, partial
from typing import (
Any,
Callable,
Expand All @@ -38,6 +38,7 @@
from sqlalchemy.types import String
from typing_extensions import TypedDict

from superset.constants import LRU_CACHE_MAX_SIZE
from superset.datasets.commands.exceptions import DatasetNotFoundError
from superset.exceptions import SupersetTemplateException
from superset.extensions import feature_flag_manager
Expand All @@ -46,7 +47,6 @@
get_user_id,
merge_extra_filters,
)
from superset.utils.memoized import memoized

if TYPE_CHECKING:
from superset.connectors.sqla.models import SqlaTable
Expand All @@ -70,7 +70,7 @@
COLLECTION_TYPES = ("list", "dict", "tuple", "set")


@memoized
@lru_cache(maxsize=LRU_CACHE_MAX_SIZE)
def context_addons() -> Dict[str, Any]:
return current_app.config.get("JINJA_CONTEXT_ADDONS", {})

Expand Down Expand Up @@ -602,7 +602,7 @@ def process_template(self, sql: str, **kwargs: Any) -> str:
}


@memoized
@lru_cache(maxsize=LRU_CACHE_MAX_SIZE)
def get_template_processors() -> Dict[str, Any]:
processors = current_app.config.get("CUSTOM_TEMPLATE_PROCESSORS", {})
for engine, processor in DEFAULT_PROCESSORS.items():
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@

from superset import db, db_engine_specs
from superset.databases.utils import make_url_safe
from superset.utils.memoized import memoized

Base = declarative_base()

Expand Down Expand Up @@ -70,7 +69,6 @@ class Slice(Base):
datasource_id = Column(Integer)


@memoized
def duration_by_name(database: Database):
return {grain.name: grain.duration for grain in database.grains()}

Expand Down
7 changes: 3 additions & 4 deletions superset/models/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
from contextlib import closing, contextmanager, nullcontext
from copy import deepcopy
from datetime import datetime
from functools import lru_cache
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, TYPE_CHECKING

import numpy
Expand Down Expand Up @@ -54,7 +55,7 @@
from sqlalchemy.sql import expression, Select

from superset import app, db_engine_specs
from superset.constants import PASSWORD_MASK
from superset.constants import LRU_CACHE_MAX_SIZE, PASSWORD_MASK
from superset.databases.utils import make_url_safe
from superset.db_engine_specs.base import MetricType, TimeGrain
from superset.extensions import (
Expand All @@ -67,7 +68,6 @@
from superset.result_set import SupersetResultSet
from superset.utils import cache as cache_util, core as utils
from superset.utils.core import get_username
from superset.utils.memoized import memoized

config = app.config
custom_password_store = config["SQLALCHEMY_CUSTOM_PASSWORD_STORE"]
Expand Down Expand Up @@ -723,7 +723,7 @@ def db_engine_spec(self) -> Type[db_engine_specs.BaseEngineSpec]:
return self.get_db_engine_spec(url)

@classmethod
@memoized
@lru_cache(maxsize=LRU_CACHE_MAX_SIZE)
def get_db_engine_spec(cls, url: URL) -> Type[db_engine_specs.BaseEngineSpec]:
backend = url.get_backend_name()
try:
Expand Down Expand Up @@ -897,7 +897,6 @@ def has_view(self, view_name: str, schema: Optional[str] = None) -> bool:
def has_view_by_name(self, view_name: str, schema: Optional[str] = None) -> bool:
return self.has_view(view_name=view_name, schema=schema)

@memoized
def get_dialect(self) -> Dialect:
sqla_url = make_url_safe(self.sqlalchemy_uri_decrypted)
return sqla_url.get_dialect()()
Expand Down
2 changes: 0 additions & 2 deletions superset/models/datasource_access_request.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@

from superset import app, db, security_manager
from superset.models.helpers import AuditMixinNullable
from superset.utils.memoized import memoized

if TYPE_CHECKING:
from superset.connectors.base.models import BaseDatasource
Expand Down Expand Up @@ -57,7 +56,6 @@ def datasource(self) -> "BaseDatasource":
return self.get_datasource

@datasource.getter # type: ignore
@memoized
def get_datasource(self) -> "BaseDatasource":
ds = db.session.query(self.cls_model).filter_by(id=self.datasource_id).first()
return ds
Expand Down
11 changes: 6 additions & 5 deletions superset/models/slice.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@
from superset.tasks.utils import get_current_user
from superset.thumbnails.digest import get_chart_digest
from superset.utils import core as utils
from superset.utils.memoized import memoized
from superset.viz import BaseViz, viz_types

if TYPE_CHECKING:
Expand Down Expand Up @@ -151,9 +150,12 @@ def clone(self) -> "Slice":

# pylint: disable=using-constant-test
@datasource.getter # type: ignore
@memoized
def get_datasource(self) -> Optional["BaseDatasource"]:
return db.session.query(self.cls_model).filter_by(id=self.datasource_id).first()
return (
db.session.query(self.cls_model)
.filter_by(id=self.datasource_id)
.one_or_none()
)

@renders("datasource_name")
def datasource_link(self) -> Optional[Markup]:
Expand Down Expand Up @@ -189,8 +191,7 @@ def datasource_edit_url(self) -> Optional[str]:

# pylint: enable=using-constant-test

@property # type: ignore
@memoized
@property
def viz(self) -> Optional[BaseViz]:
form_data = json.loads(self.params)
viz_class = viz_types.get(self.viz_type)
Expand Down
6 changes: 3 additions & 3 deletions superset/utils/date_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import logging
import re
from datetime import datetime, timedelta
from functools import lru_cache
from time import struct_time
from typing import Dict, List, Optional, Tuple

Expand Down Expand Up @@ -45,8 +46,7 @@
TimeRangeAmbiguousError,
TimeRangeParseFailError,
)
from superset.constants import NO_TIME_RANGE
from superset.utils.memoized import memoized
from superset.constants import LRU_CACHE_MAX_SIZE, NO_TIME_RANGE

ParserElement.enablePackrat()

Expand Down Expand Up @@ -394,7 +394,7 @@ def eval(self) -> datetime:
)


@memoized
@lru_cache(maxsize=LRU_CACHE_MAX_SIZE)
def datetime_parser() -> ParseResults: # pylint: disable=too-many-locals
( # pylint: disable=invalid-name
DATETIME,
Expand Down
81 changes: 0 additions & 81 deletions superset/utils/memoized.py

This file was deleted.

1 change: 1 addition & 0 deletions tests/integration_tests/utils_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -991,6 +991,7 @@ def test_log_this(self) -> None:
slc = self.get_slice("Girls", db.session)
dashboard_id = 1

assert slc.viz is not None
resp = self.get_json_resp(
f"/superset/explore_json/{slc.datasource_type}/{slc.datasource_id}/"
+ f'?form_data={{"slice_id": {slc.id}}}&dashboard_id={dashboard_id}',
Expand Down
Loading

0 comments on commit ef648ae

Please sign in to comment.