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

superset/commands/dashboard/permalink/create.py

@@ -18,6 +18,7 @@ import logging
 
 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 @@ class CreateDashboardPermalinkCommand(BaseDashboardPermalinkCommand):
                 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 typing import Any, Optional
 
 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 @@ class CreateExplorePermalinkCommand(BaseExplorePermalinkCommand):
             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 @@ class CreateKeyValueCommand(BaseCommand):
             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 @@ class DeleteKeyValueCommand(BaseCommand):
         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 @@ class DeleteExpiredKeyValueCommand(BaseCommand):
             )
             .delete()
         )
-        db.session.commit()
+        db.session.flush()

superset/commands/key_value/update.py

@@ -84,7 +84,7 @@ class UpdateKeyValueCommand(BaseCommand):
             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 @@ class UpsertKeyValueCommand(BaseCommand):
             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 uuid import UUID, uuid3
 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,6 +95,7 @@ class SupersetMetastoreCache(BaseCache):
             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:
@@ -109,6 +111,7 @@ class SupersetMetastoreCache(BaseCache):
                 key=self.get_key(key),
                 expires_on=self._get_expiry(timeout),
             ).run()
+            db.session.commit()
             return True
         except KeyValueCreateFailedError:
             return False
@@ -133,4 +136,6 @@ class SupersetMetastoreCache(BaseCache):
         # 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/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 @@ def set_shared_value(key: SharedKey, value: Any) -> None:
         key=uuid_key,
         codec=CODEC,
     ).run()
+    db.session.commit()
 
 
 def get_permalink_salt(key: SharedKey) -> str:

superset/utils/lock.py

@@ -24,6 +24,7 @@ from contextlib import contextmanager
 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 JsonKeyValueCodec, KeyValueResource
@@ -92,10 +93,12 @@ def KeyValueDistributedLock(  # pylint: disable=invalid-name
             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/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)