use type alias for key

apache · Jun 30, 2024 · 4e7c75f · 4e7c75f
1 parent 515dc80
commit 4e7c75f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 15 deletions.
diff --git a/superset/daos/key_value.py b/superset/daos/key_value.py
@@ -30,7 +30,7 @@
     KeyValueUpdateFailedError,
 )
 from superset.key_value.models import KeyValueEntry
-from superset.key_value.types import KeyValueCodec, KeyValueResource
+from superset.key_value.types import Key, KeyValueCodec, KeyValueResource
 from superset.key_value.utils import get_filter
 from superset.utils.core import get_user_id
 
@@ -41,22 +41,16 @@ class KeyValueDAO(BaseDAO[KeyValueEntry]):
     @staticmethod
     def get_entry(
         resource: KeyValueResource,
-        key: int | UUID,
+        key: Key,
     ) -> KeyValueEntry | None:
-        """
-
-        :param resource:
-        :param key:
-        :return:
-        """
         filter_ = get_filter(resource, key)
         return db.session.query(KeyValueEntry).filter_by(**filter_).first()
 
     @classmethod
     def get_value(
         cls,
         resource: KeyValueResource,
-        key: int | UUID,
+        key: Key,
         codec: KeyValueCodec,
     ) -> Any:
         entry = cls.get_entry(resource, key)
@@ -66,7 +60,7 @@ def get_value(
         return codec.decode(entry.value)
 
     @staticmethod
-    def delete_entry(resource: KeyValueResource, key: int | UUID) -> bool:
+    def delete_entry(resource: KeyValueResource, key: Key) -> bool:
         if entry := KeyValueDAO.get_entry(resource, key):
             db.session.delete(entry)
             return True
@@ -91,7 +85,7 @@ def create_entry(
         resource: KeyValueResource,
         value: Any,
         codec: KeyValueCodec,
-        key: int | UUID | None = None,
+        key: Key | None = None,
         expires_on: datetime | None = None,
     ) -> KeyValueEntry:
         try:
@@ -121,7 +115,7 @@ def upsert_entry(
         resource: KeyValueResource,
         value: Any,
         codec: KeyValueCodec,
-        key: int | UUID,
+        key: Key,
         expires_on: datetime | None = None,
     ) -> KeyValueEntry:
         if entry := KeyValueDAO.get_entry(resource, key):
@@ -138,7 +132,7 @@ def update_entry(
         resource: KeyValueResource,
         value: Any,
         codec: KeyValueCodec,
-        key: int | UUID,
+        key: Key,
         expires_on: datetime | None = None,
     ) -> KeyValueEntry:
         if entry := KeyValueDAO.get_entry(resource, key):

diff --git a/superset/key_value/types.py b/superset/key_value/types.py
@@ -30,6 +30,8 @@
 )
 from superset.utils.backports import StrEnum
 
+Key = int | UUID
+
 
 class KeyValueFilter(TypedDict, total=False):
     resource: str

diff --git a/superset/key_value/utils.py b/superset/key_value/utils.py
@@ -25,7 +25,7 @@
 from flask_babel import gettext as _
 
 from superset.key_value.exceptions import KeyValueParseKeyError
-from superset.key_value.types import KeyValueFilter, KeyValueResource
+from superset.key_value.types import Key, KeyValueFilter, KeyValueResource
 from superset.utils.json import json_dumps_w_dates
 
 HASHIDS_MIN_LENGTH = 11
@@ -35,7 +35,7 @@ def random_key() -> str:
     return token_urlsafe(48)
 
 
-def get_filter(resource: KeyValueResource, key: int | UUID) -> KeyValueFilter:
+def get_filter(resource: KeyValueResource, key: Key) -> KeyValueFilter:
     try:
         filter_: KeyValueFilter = {"resource": resource.value}
         if isinstance(key, UUID):