fix(key-value): use flush instead of commit

superset/commands/dashboard/permalink/create.py

@@ -18,6 +18,7 @@
 
 from sqlalchemy.exc import SQLAlchemyError
 
+from superset import db
 from superset.commands.dashboard.permalink.base import BaseDashboardPermalinkCommand
 from superset.commands.key_value.upsert import UpsertKeyValueCommand
 from superset.daos.dashboard import DashboardDAO
@@ -62,6 +63,7 @@
                 codec=self.codec,
             ).run()
             assert key.id  # for type checks
+            db.session.commit()
             return encode_permalink_key(key=key.id, salt=self.salt)
         except KeyValueCodecEncodeException as ex:
             raise DashboardPermalinkCreateFailedError(str(ex)) from ex

superset/commands/explore/permalink/create.py

@@ -19,6 +19,7 @@
 
 from sqlalchemy.exc import SQLAlchemyError
 
+from superset import db
 from superset.commands.explore.permalink.base import BaseExplorePermalinkCommand
 from superset.commands.key_value.create import CreateKeyValueCommand
 from superset.explore.permalink.exceptions import ExplorePermalinkCreateFailedError
@@ -58,6 +59,7 @@
             key = command.run()
             if key.id is None:
                 raise ExplorePermalinkCreateFailedError("Unexpected missing key id")
+            db.session.commit()
             return encode_permalink_key(key=key.id, salt=self.salt)
         except KeyValueCodecEncodeException as ex:
             raise ExplorePermalinkCreateFailedError(str(ex)) from ex

superset/commands/key_value/create.py

@@ -99,5 +99,5 @@ def create(self) -> Key:
             except ValueError as ex:
                 raise KeyValueCreateFailedError() from ex
         db.session.add(entry)
-        db.session.commit()
+        db.session.flush()
         return Key(id=entry.id, uuid=entry.uuid)

superset/commands/key_value/delete.py

@@ -59,6 +59,6 @@ def delete(self) -> bool:
         filter_ = get_filter(self.resource, self.key)
         if entry := db.session.query(KeyValueEntry).filter_by(**filter_).first():
             db.session.delete(entry)
-            db.session.commit()
+            db.session.flush()
             return True
         return False

superset/commands/key_value/delete_expired.py

@@ -62,4 +62,4 @@ def delete_expired(self) -> None:
             )
             .delete()
         )
-        db.session.commit()
+        db.session.flush()

superset/commands/key_value/get.py

@@ -16,7 +16,6 @@
 # under the License.
 
 import logging
-from datetime import datetime
 from typing import Any, Optional, Union
 from uuid import UUID
 
@@ -67,6 +66,6 @@
     def get(self) -> Optional[Any]:
         filter_ = get_filter(self.resource, self.key)
         entry = db.session.query(KeyValueEntry).filter_by(**filter_).first()
-        if entry and (entry.expires_on is None or entry.expires_on > datetime.now()):
+        if entry and not entry.is_expired():
             return self.codec.decode(entry.value)
         return None

superset/commands/key_value/update.py

@@ -84,7 +84,7 @@
             entry.expires_on = self.expires_on
             entry.changed_on = datetime.now()
             entry.changed_by_fk = get_user_id()
-            db.session.commit()
+            db.session.flush()
             return Key(id=entry.id, uuid=entry.uuid)
 
         return None

superset/commands/key_value/upsert.py

@@ -88,7 +88,7 @@
             entry.expires_on = self.expires_on
             entry.changed_on = datetime.now()
             entry.changed_by_fk = get_user_id()
-            db.session.commit()
+            db.session.flush()
             return Key(entry.id, entry.uuid)
 
         return CreateKeyValueCommand(

superset/extensions/metastore_cache.py

@@ -22,6 +22,7 @@
 from flask import current_app, Flask, has_app_context
 from flask_caching import BaseCache
 
+from superset import db
 from superset.key_value.exceptions import KeyValueCreateFailedError
 from superset.key_value.types import (
     KeyValueCodec,
@@ -94,21 +95,23 @@
             codec=self.codec,
             expires_on=self._get_expiry(timeout),
         ).run()
+        db.session.commit()
         return True
 
     def add(self, key: str, value: Any, timeout: Optional[int] = None) -> bool:
         # pylint: disable=import-outside-toplevel
         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()
+            db.session.commit()
             return True
         except KeyValueCreateFailedError:
             return False
@@ -133,4 +136,6 @@
         # pylint: disable=import-outside-toplevel
         from superset.commands.key_value.delete import DeleteKeyValueCommand
 
-        return DeleteKeyValueCommand(resource=RESOURCE, key=self.get_key(key)).run()
+        ret = DeleteKeyValueCommand(resource=RESOURCE, key=self.get_key(key)).run()
+        db.session.commit()
+        return ret

superset/key_value/models.py

@@ -14,6 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from datetime import datetime
+
 from flask_appbuilder import Model
 from sqlalchemy import Column, DateTime, ForeignKey, Integer, LargeBinary, String
 from sqlalchemy.orm import relationship
@@ -38,3 +40,6 @@ class KeyValueEntry(AuditMixinNullable, ImportExportMixin, Model):
     changed_by_fk = Column(Integer, ForeignKey("ab_user.id"), nullable=True)
     created_by = relationship(security_manager.user_model, foreign_keys=[created_by_fk])
     changed_by = relationship(security_manager.user_model, foreign_keys=[changed_by_fk])
+
+    def is_expired(self) -> bool:
+        return self.expires_on is not None and self.expires_on <= datetime.now()

superset/key_value/shared_entries.py

@@ -18,6 +18,7 @@
 from typing import Any, Optional
 from uuid import uuid3
 
+from superset import db
 from superset.key_value.types import JsonKeyValueCodec, KeyValueResource, SharedKey
 from superset.key_value.utils import get_uuid_namespace, random_key
 
@@ -45,6 +46,7 @@
         key=uuid_key,
         codec=CODEC,
     ).run()
+    db.session.commit()
 
 
 def get_permalink_salt(key: SharedKey) -> str:

superset/utils/lock.py

@@ -24,9 +24,10 @@
 from datetime import datetime, timedelta
 from typing import Any, cast, TypeVar, Union
 
+from superset import db
 from superset.exceptions import CreateKeyValueDistributedLockFailedException
 from superset.key_value.exceptions import KeyValueCreateFailedError
-from superset.key_value.types import KeyValueResource, PickleKeyValueCodec
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
 from superset.utils import json
 
 LOCK_EXPIRATION = timedelta(seconds=30)
@@ -53,6 +54,10 @@ def sort(obj: T) -> T:
     return json.dumps(params)
 
 
+def get_key(namespace: str, **kwargs: Any) -> uuid.UUID:
+    return uuid.uuid5(uuid.uuid5(uuid.NAMESPACE_DNS, namespace), serialize(kwargs))
+
+
 @contextmanager
 def KeyValueDistributedLock(  # pylint: disable=invalid-name
     namespace: str,
@@ -77,21 +82,23 @@ def KeyValueDistributedLock(  # pylint: disable=invalid-name
     from superset.commands.key_value.delete import DeleteKeyValueCommand
     from superset.commands.key_value.delete_expired import DeleteExpiredKeyValueCommand
 
-    key = uuid.uuid5(uuid.uuid5(uuid.NAMESPACE_DNS, namespace), serialize(kwargs))
+    key = get_key(namespace, **kwargs)
     logger.debug("Acquiring lock on namespace %s for key %s", namespace, key)
     try:
         DeleteExpiredKeyValueCommand(resource=KeyValueResource.LOCK).run()
         CreateKeyValueCommand(
             resource=KeyValueResource.LOCK,
-            codec=PickleKeyValueCodec(),
+            codec=JsonKeyValueCodec(),
             key=key,
             value=True,
             expires_on=datetime.now() + LOCK_EXPIRATION,
         ).run()
+        db.session.commit()
 
         yield key
 
         DeleteKeyValueCommand(resource=KeyValueResource.LOCK, key=key).run()
+        db.session.commit()
         logger.debug("Removed lock on namespace %s for key %s", namespace, key)
     except KeyValueCreateFailedError as ex:
         raise CreateKeyValueDistributedLockFailedException(

tests/integration_tests/extensions/metastore_cache_test.py

@@ -77,15 +77,24 @@ 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.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.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
+        cache.add(FIRST_KEY, SECOND_VALUE, int(delta.total_seconds()))
+        assert cache.get(FIRST_KEY) == SECOND_VALUE
 
 
 @pytest.mark.parametrize(

tests/integration_tests/key_value/commands/delete_test.py

@@ -49,7 +49,7 @@ def key_value_entry() -> KeyValueEntry:
         value=bytes(json.dumps(JSON_VALUE), encoding="utf-8"),
     )
     db.session.add(entry)
-    db.session.commit()
+    db.session.flush()
     return entry
 
 
@@ -61,6 +61,7 @@ def test_delete_id_entry(
     from superset.commands.key_value.delete import DeleteKeyValueCommand
 
     assert DeleteKeyValueCommand(resource=RESOURCE, key=ID_KEY).run() is True
+    db.session.commit()
 
 
 def test_delete_uuid_entry(
@@ -71,12 +72,12 @@ def test_delete_uuid_entry(
     from superset.commands.key_value.delete import DeleteKeyValueCommand
 
     assert DeleteKeyValueCommand(resource=RESOURCE, key=UUID_KEY).run() is True
+    db.session.commit()
 
 
 def test_delete_entry_missing(
     app_context: AppContext,
     admin: User,  # noqa: F811
-    key_value_entry: KeyValueEntry,
 ) -> None:
     from superset.commands.key_value.delete import DeleteKeyValueCommand
 

tests/integration_tests/key_value/commands/fixtures.py

@@ -56,7 +56,7 @@ def key_value_entry() -> Generator[KeyValueEntry, None, None]:
         value=bytes(json.dumps(JSON_VALUE), encoding="utf-8"),
     )
     db.session.add(entry)
-    db.session.commit()
+    db.session.flush()
     yield entry
     db.session.delete(entry)
     db.session.commit()

tests/integration_tests/key_value/commands/get_test.py

@@ -76,7 +76,7 @@ def test_get_expired_entry(app_context: AppContext) -> None:
         expires_on=datetime.now() - timedelta(days=1),
     )
     db.session.add(entry)
-    db.session.commit()
+    db.session.flush()
     value = GetKeyValueCommand(resource=RESOURCE, key=ID_KEY, codec=JSON_CODEC).run()
     assert value is None
     db.session.delete(entry)
@@ -96,7 +96,7 @@ def test_get_future_expiring_entry(app_context: AppContext) -> None:
         expires_on=datetime.now() + timedelta(days=1),
     )
     db.session.add(entry)
-    db.session.commit()
+    db.session.flush()
     value = GetKeyValueCommand(resource=RESOURCE, key=id_, codec=JSON_CODEC).run()
     assert value == JSON_VALUE
     db.session.delete(entry)

tests/integration_tests/key_value/commands/upsert_test.py

@@ -98,4 +98,3 @@ def test_upsert_missing_entry(app_context: AppContext, admin: User) -> None:  #
     assert key is not None
     assert key.id == 456
     db.session.query(KeyValueEntry).filter_by(id=456).delete()
-    db.session.commit()