fix(metastore-cache): prune before add (#29301)

apache · Jun 20, 2024 · 172ddb4 · 172ddb4
1 parent 99fc04b
commit 172ddb4
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/superset/extensions/metastore_cache.py b/superset/extensions/metastore_cache.py
@@ -101,14 +101,14 @@ def add(self, key: str, value: Any, timeout: Optional[int] = None) -> bool:
         from superset.commands.key_value.create import CreateKeyValueCommand
 
         try:
+            self._prune()
             CreateKeyValueCommand(
                 resource=RESOURCE,
                 value=value,
                 codec=self.codec,
                 key=self.get_key(key),
                 expires_on=self._get_expiry(timeout),
             ).run()
-            self._prune()
             return True
         except KeyValueCreateFailedError:
             return False

diff --git a/tests/integration_tests/extensions/metastore_cache_test.py b/tests/integration_tests/extensions/metastore_cache_test.py
@@ -77,16 +77,29 @@ def test_caching_flow(app_context: AppContext, cache: SupersetMetastoreCache) ->
 def test_expiry(app_context: AppContext, cache: SupersetMetastoreCache) -> None:
     delta = timedelta(days=90)
     dttm = datetime(2022, 3, 18, 0, 0, 0)
+
+    # 1. initialize cached values, ensure they're found
     with freeze_time(dttm):
-        cache.set(FIRST_KEY, FIRST_KEY_INITIAL_VALUE, int(delta.total_seconds()))
+        assert (
+            cache.set(FIRST_KEY, FIRST_KEY_INITIAL_VALUE, int(delta.total_seconds()))
+            is True
+        )
         assert cache.get(FIRST_KEY) == FIRST_KEY_INITIAL_VALUE
+
+    # 2. ensure cached values are available a moment before expiration
     with freeze_time(dttm + delta - timedelta(seconds=1)):
-        assert cache.has(FIRST_KEY)
+        assert cache.has(FIRST_KEY) is True
         assert cache.get(FIRST_KEY) == FIRST_KEY_INITIAL_VALUE
+
+    # 3. ensure cached entries expire
     with freeze_time(dttm + delta + timedelta(seconds=1)):
         assert cache.has(FIRST_KEY) is False
         assert cache.get(FIRST_KEY) is None
 
+        # adding a value with the same key as an expired entry works
+        assert cache.add(FIRST_KEY, SECOND_VALUE, int(delta.total_seconds())) is True
+        assert cache.get(FIRST_KEY) == SECOND_VALUE
+
 
 @pytest.mark.parametrize(
     "input_,codec,expected_result",