From 7f3c8efab001a01e0a7c146f14f04be855bf82de Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Tue, 2 Jul 2024 15:28:42 +0300
Subject: [PATCH] fix(metastore-cache): import dao in methods (#29451)

---
 superset/extensions/metastore_cache.py              | 13 ++++++++++++-
 .../extensions/metastore_cache_test.py              |  8 ++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/superset/extensions/metastore_cache.py b/superset/extensions/metastore_cache.py
index 1195bd8edf65d..197dac9386d3f 100644
--- a/superset/extensions/metastore_cache.py
+++ b/superset/extensions/metastore_cache.py
@@ -24,7 +24,6 @@
 from sqlalchemy.exc import SQLAlchemyError
 
 from superset import db
-from superset.daos.key_value import KeyValueDAO
 from superset.key_value.exceptions import KeyValueCreateFailedError
 from superset.key_value.types import (
     KeyValueCodec,
@@ -79,6 +78,9 @@ def _get_expiry(self, timeout: Optional[int]) -> Optional[datetime]:
         return None
 
     def set(self, key: str, value: Any, timeout: Optional[int] = None) -> bool:
+        # pylint: disable=import-outside-toplevel
+        from superset.daos.key_value import KeyValueDAO
+
         KeyValueDAO.upsert_entry(
             resource=RESOURCE,
             key=self.get_key(key),
@@ -90,6 +92,9 @@ def set(self, key: str, value: Any, timeout: Optional[int] = None) -> bool:
         return True
 
     def add(self, key: str, value: Any, timeout: Optional[int] = None) -> bool:
+        # pylint: disable=import-outside-toplevel
+        from superset.daos.key_value import KeyValueDAO
+
         try:
             KeyValueDAO.delete_expired_entries(RESOURCE)
             KeyValueDAO.create_entry(
@@ -106,6 +111,9 @@ def add(self, key: str, value: Any, timeout: Optional[int] = None) -> bool:
             return False
 
     def get(self, key: str) -> Any:
+        # pylint: disable=import-outside-toplevel
+        from superset.daos.key_value import KeyValueDAO
+
         return KeyValueDAO.get_value(RESOURCE, self.get_key(key), self.codec)
 
     def has(self, key: str) -> bool:
@@ -116,4 +124,7 @@ def has(self, key: str) -> bool:
 
     @transaction()
     def delete(self, key: str) -> Any:
+        # pylint: disable=import-outside-toplevel
+        from superset.daos.key_value import KeyValueDAO
+
         return KeyValueDAO.delete_entry(RESOURCE, self.get_key(key))
diff --git a/tests/integration_tests/extensions/metastore_cache_test.py b/tests/integration_tests/extensions/metastore_cache_test.py
index 238e8fd46a506..401c2f988bed6 100644
--- a/tests/integration_tests/extensions/metastore_cache_test.py
+++ b/tests/integration_tests/extensions/metastore_cache_test.py
@@ -18,13 +18,14 @@
 
 from contextlib import nullcontext
 from datetime import datetime, timedelta
-from typing import Any, TYPE_CHECKING
+from typing import Any
 from uuid import UUID
 
 import pytest
 from flask.ctx import AppContext
 from freezegun import freeze_time
 
+from superset.extensions.metastore_cache import SupersetMetastoreCache
 from superset.key_value.exceptions import KeyValueCodecEncodeException
 from superset.key_value.types import (
     JsonKeyValueCodec,
@@ -32,9 +33,6 @@
     PickleKeyValueCodec,
 )
 
-if TYPE_CHECKING:
-    from superset.extensions.metastore_cache import SupersetMetastoreCache
-
 NAMESPACE = UUID("ee173d1b-ccf3-40aa-941c-985c15224496")
 
 FIRST_KEY = "foo"
@@ -47,8 +45,6 @@
 
 @pytest.fixture
 def cache() -> SupersetMetastoreCache:
-    from superset.extensions.metastore_cache import SupersetMetastoreCache
-
     return SupersetMetastoreCache(
         namespace=NAMESPACE,
         default_timeout=600,