chore(key-value): convert command to dao

superset/commands/dashboard/permalink/create.py

@@ -20,8 +20,8 @@
 
 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
+from superset.daos.key_value import KeyValueDAO
 from superset.dashboards.permalink.exceptions import DashboardPermalinkCreateFailedError
 from superset.dashboards.permalink.types import DashboardPermalinkState
 from superset.key_value.exceptions import KeyValueCodecEncodeException
@@ -56,15 +56,16 @@ def run(self) -> str:
                 "state": self.state,
             }
             user_id = get_user_id()
-            key = UpsertKeyValueCommand(
+            entry = KeyValueDAO.upsert_entry(
                 resource=self.resource,
                 key=get_deterministic_uuid(self.salt, (user_id, value)),
                 value=value,
                 codec=self.codec,
-            ).run()
-            assert key.id  # for type checks
+            )
+            db.session.flush()
+            assert entry.id  # for type checks
             db.session.commit()
-            return encode_permalink_key(key=key.id, salt=self.salt)
+            return encode_permalink_key(key=entry.id, salt=self.salt)
         except KeyValueCodecEncodeException as ex:
             raise DashboardPermalinkCreateFailedError(str(ex)) from ex
         except SQLAlchemyError as ex:

superset/commands/dashboard/permalink/get.py

@@ -21,8 +21,8 @@
 
 from superset.commands.dashboard.exceptions import DashboardNotFoundError
 from superset.commands.dashboard.permalink.base import BaseDashboardPermalinkCommand
-from superset.commands.key_value.get import GetKeyValueCommand
 from superset.daos.dashboard import DashboardDAO
+from superset.daos.key_value import KeyValueDAO
 from superset.dashboards.permalink.exceptions import DashboardPermalinkGetFailedError
 from superset.dashboards.permalink.types import DashboardPermalinkValue
 from superset.key_value.exceptions import (
@@ -43,12 +43,7 @@ def run(self) -> Optional[DashboardPermalinkValue]:
         self.validate()
         try:
             key = decode_permalink_id(self.key, salt=self.salt)
-            command = GetKeyValueCommand(
-                resource=self.resource,
-                key=key,
-                codec=self.codec,
-            )
-            value: Optional[DashboardPermalinkValue] = command.run()
+            value = KeyValueDAO.get_value(self.resource, key, self.codec)
             if value:
                 DashboardDAO.get_by_id_or_slug(value["dashboardId"])
                 return value

tests/integration_tests/key_value/__init__.py → superset/commands/distributed_lock/base.py

@@ -14,3 +14,29 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+import logging
+import uuid
+from typing import Any
+
+from flask import current_app
+
+from superset.commands.base import BaseCommand
+from superset.distributed_lock.utils import get_key
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
+
+logger = logging.getLogger(__name__)
+stats_logger = current_app.config["STATS_LOGGER"]
+
+CODEC = JsonKeyValueCodec()
+RESOURCE = KeyValueResource.LOCK
+
+
+class BaseDistributedLockCommand(BaseCommand):
+    key: uuid.UUID
+
+    def __init__(self, namespace: str, params: dict[str, Any] | None = None):
+        self.key = get_key(namespace, **(params or {}))
+
+    def validate(self) -> None:
+        pass

superset/commands/distributed_lock/create.py

@@ -0,0 +1,59 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import logging
+import uuid
+from datetime import datetime, timedelta
+
+from flask import current_app
+from sqlalchemy.exc import IntegrityError
+
+from superset import db
+from superset.commands.distributed_lock.base import BaseDistributedLockCommand
+from superset.daos.key_value import KeyValueDAO
+from superset.exceptions import CreateKeyValueDistributedLockFailedException
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
+
+logger = logging.getLogger(__name__)
+stats_logger = current_app.config["STATS_LOGGER"]
+
+
+class CreateDistributedLock(BaseDistributedLockCommand):
+    key: uuid.UUID
+    lock_expiration = timedelta(seconds=30)
+    resource = KeyValueResource.LOCK
+    codec = JsonKeyValueCodec()
+
+    def validate(self) -> None:
+        pass
+
+    def run(self) -> None:
+        try:
+            KeyValueDAO.delete_expired_entries(self.resource)
+            KeyValueDAO.create_entry(
+                resource=KeyValueResource.LOCK,
+                value=True,
+                codec=self.codec,
+                key=self.key,
+                expires_on=datetime.now() + self.lock_expiration,
+            )
+            db.session.commit()
+        except IntegrityError as ex:
+            db.session.rollback()
+            raise CreateKeyValueDistributedLockFailedException(
+                "Lock already taken"
+            ) from ex

superset/commands/key_value/__init__.py → superset/commands/distributed_lock/delete.py

@@ -14,3 +14,29 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+import logging
+import uuid
+
+from flask import current_app
+
+from superset import db
+from superset.commands.distributed_lock.base import BaseDistributedLockCommand
+from superset.daos.key_value import KeyValueDAO
+from superset.key_value.types import KeyValueResource
+
+logger = logging.getLogger(__name__)
+stats_logger = current_app.config["STATS_LOGGER"]
+
+RESOURCE = KeyValueResource.LOCK
+
+
+class DeleteDistributedLock(BaseDistributedLockCommand):
+    key: uuid.UUID
+
+    def validate(self) -> None:
+        pass
+
+    def run(self) -> None:
+        KeyValueDAO.delete_entry(RESOURCE, self.key)
+        db.session.commit()

superset/commands/explore/permalink/create.py

@@ -21,7 +21,7 @@
 
 from superset import db
 from superset.commands.explore.permalink.base import BaseExplorePermalinkCommand
-from superset.commands.key_value.create import CreateKeyValueCommand
+from superset.daos.key_value import KeyValueDAO
 from superset.explore.permalink.exceptions import ExplorePermalinkCreateFailedError
 from superset.explore.utils import check_access as check_chart_access
 from superset.key_value.exceptions import KeyValueCodecEncodeException
@@ -51,16 +51,13 @@ def run(self) -> str:
                 "datasource": self.datasource,
                 "state": self.state,
             }
-            command = CreateKeyValueCommand(
-                resource=self.resource,
-                value=value,
-                codec=self.codec,
-            )
-            key = command.run()
-            if key.id is None:
+            entry = KeyValueDAO.create_entry(self.resource, value, self.codec)
+            db.session.flush()
+            key = entry.id
+            if key is None:
                 raise ExplorePermalinkCreateFailedError("Unexpected missing key id")
             db.session.commit()
-            return encode_permalink_key(key=key.id, salt=self.salt)
+            return encode_permalink_key(key=key, salt=self.salt)
         except KeyValueCodecEncodeException as ex:
             raise ExplorePermalinkCreateFailedError(str(ex)) from ex
         except SQLAlchemyError as ex:

superset/commands/explore/permalink/get.py

@@ -21,7 +21,7 @@
 
 from superset.commands.dataset.exceptions import DatasetNotFoundError
 from superset.commands.explore.permalink.base import BaseExplorePermalinkCommand
-from superset.commands.key_value.get import GetKeyValueCommand
+from superset.daos.key_value import KeyValueDAO
 from superset.explore.permalink.exceptions import ExplorePermalinkGetFailedError
 from superset.explore.permalink.types import ExplorePermalinkValue
 from superset.explore.utils import check_access as check_chart_access
@@ -44,11 +44,7 @@ def run(self) -> Optional[ExplorePermalinkValue]:
         self.validate()
         try:
             key = decode_permalink_id(self.key, salt=self.salt)
-            value: Optional[ExplorePermalinkValue] = GetKeyValueCommand(
-                resource=self.resource,
-                key=key,
-                codec=self.codec,
-            ).run()
+            value = KeyValueDAO.get_value(self.resource, key, self.codec)
             if value:
                 chart_id: Optional[int] = value.get("chartId")
                 # keep this backward compatible for old permalinks