refactor expiry logic

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 @@ def run(self) -> str:
                 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 @@ def run(self) -> str:
             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/get.py

@@ -16,12 +16,10 @@
 # under the License.
 
 import logging
-from datetime import datetime
 from typing import Any, Optional, Union
 from uuid import UUID
 
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.orm import Session
 
 from superset import db
 from superset.commands.base import BaseCommand
@@ -37,28 +35,24 @@ class GetKeyValueCommand(BaseCommand):
     resource: KeyValueResource
     key: Union[int, UUID]
     codec: KeyValueCodec
-    session: Session  # pylint: disable=disallowed-name
 
     def __init__(
         self,
         resource: KeyValueResource,
         key: Union[int, UUID],
         codec: KeyValueCodec,
-        session: Optional[Session] = None,  # pylint: disable=disallowed-name
     ):
         """
         Retrieve a key value entry
 
         :param resource: the resource (dashboard, chart etc)
         :param key: the key to retrieve
         :param codec: codec used to decode the value
-        :param session: session to bind to (defaults to superset.db.session)
         :return: the value associated with the key if present
         """
         self.resource = resource
         self.key = key
         self.codec = codec
-        self.session = session or db.session  # pylint: disable=disallowed-name
 
     def run(self) -> Any:
         try:
@@ -71,7 +65,7 @@ def validate(self) -> None:
 
     def get(self) -> Optional[Any]:
         filter_ = get_filter(self.resource, self.key)
-        entry = self.session.query(KeyValueEntry).filter_by(**filter_).first()
-        if entry and (entry.expires_on is None or entry.expires_on > datetime.now()):
+        entry = db.session.query(KeyValueEntry).filter_by(**filter_).first()
+        if entry and not entry.is_expired():
             return self.codec.decode(entry.value)
         return None

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

tests/unit_tests/utils/lock_tests.py

@@ -24,22 +24,22 @@
 from freezegun import freeze_time
 from sqlalchemy.orm import Session, sessionmaker
 
-from superset.commands.key_value.get import GetKeyValueCommand
 from superset.exceptions import CreateKeyValueDistributedLockFailedException
-from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
+from superset.key_value.types import JsonKeyValueCodec
 from superset.utils.lock import get_key, KeyValueDistributedLock
 
 MAIN_KEY = get_key("ns", a=1, b=2)
 OTHER_KEY = get_key("ns2", a=1, b=2)
 
 
 def _get_lock(key: UUID, session: Session) -> Any:
-    return GetKeyValueCommand(
-        resource=KeyValueResource.LOCK,
-        key=key,
-        codec=JsonKeyValueCodec(),
-        session=session,
-    ).run()
+    from superset.key_value.models import KeyValueEntry
+
+    entry = session.query(KeyValueEntry).filter_by(uuid=key).first()
+    if entry is None or entry.is_expired():
+        return None
+
+    return JsonKeyValueCodec().decode(entry.value)
 
 
 def _get_session() -> Session: