fix: memoized decorator memory leak (#23139)

(cherry picked from commit 79274eb)

superset/connectors/sqla/utils.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import logging
+from functools import lru_cache
 from typing import (
     Any,
     Callable,
@@ -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,
@@ -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
@@ -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
 

superset/constants.py

@@ -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
     """

superset/jinja_context.py

@@ -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,
@@ -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
@@ -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
@@ -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", {})
 
@@ -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():

superset/migrations/versions/2020-04-29_09-24_620241d1153f_update_time_grain_sqla.py

@@ -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()
 
@@ -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()}
 

superset/models/core.py

@@ -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
@@ -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 (
@@ -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"]
@@ -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:
@@ -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()()

superset/models/datasource_access_request.py

@@ -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
@@ -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

superset/models/slice.py

@@ -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:
@@ -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]:
@@ -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)

superset/utils/date_parser.py

@@ -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
 
@@ -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()
 
@@ -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,

tests/integration_tests/utils_tests.py

@@ -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}',