chore: Move constants and types from deferred deletion

Simple update for storage, now using contextmanager instead of an intermediary class for "as dict"

server/src/uds/core/consts/defered_deletion.py → server/src/uds/core/consts/deferred_deletion.py

@@ -48,8 +48,8 @@
 CHECK_INTERVAL: typing.Final[datetime.timedelta] = datetime.timedelta(seconds=11)  # Check interval
 FATAL_ERROR_INTERVAL_MULTIPLIER: typing.Final[int] = 2  # Multiplier for fatal errors
 
-TO_STOP_GROUP: typing.Final[str] = 'to_stop'
-STOPPING_GROUP: typing.Final[str] = 'stopping'
-TO_DELETE_GROUP: typing.Final[str] = 'to_delete'
-DELETING_GROUP: typing.Final[str] = 'deleting'
+# TO_STOP_GROUP: typing.Final[str] = 'to_stop'
+# STOPPING_GROUP: typing.Final[str] = 'stopping'
+# TO_DELETE_GROUP: typing.Final[str] = 'to_delete'
+# DELETING_GROUP: typing.Final[str] = 'deleting'
 

server/src/uds/core/managers/servers.py

@@ -28,6 +28,7 @@
 """
 Author: Adolfo Gómez, dkmaster at dkmon dot com
 """
+import contextlib
 import datetime
 import logging
 import typing
@@ -41,7 +42,7 @@
 from uds.core import exceptions, types
 from uds.core.util import model as model_utils
 from uds.core.util import singleton
-from uds.core.util.storage import StorageAccess, Storage
+from uds.core.util.storage import StorageAsDict, Storage
 
 from .servers_api import events, requester
 
@@ -65,11 +66,13 @@ class ServerManager(metaclass=singleton.Singleton):
     def manager() -> 'ServerManager':
         return ServerManager()  # Singleton pattern will return always the same instance
 
-    def counter_storage(self) -> 'StorageAccess':
+    @contextlib.contextmanager
+    def counter_storage(self) -> typing.Iterator[StorageAsDict]:
         # If counters are too old, restart them
         if datetime.datetime.now() - self.last_counters_clean > self.MAX_COUNTERS_AGE:
             self.clear_unmanaged_usage()
-        return Storage(self.STORAGE_NAME).as_dict(atomic=True, group='counters')
+        with Storage(self.STORAGE_NAME).as_dict(atomic=True, group='counters') as storage:
+            yield storage
 
     def property_name(self, user: typing.Optional[typing.Union[str, 'models.User']]) -> str:
         """Returns the property name for a user"""
@@ -155,7 +158,9 @@ def _real_weight(stats: 'types.servers.ServerStats') -> float:
             # To values over threshold, we will add 1, so they are always worse than any value under threshold
             # No matter if over threshold is overcalculed, it will be always worse than any value under threshold
             # and all values over threshold will be affected in the same way
-            return weight_threshold - stats.weight() if stats.weight() < weight_threshold else 1 + stats.weight()
+            return (
+                weight_threshold - stats.weight() if stats.weight() < weight_threshold else 1 + stats.weight()
+            )
 
         # Now, cachedStats has a list of tuples (stats, server), use it to find the best server
         for stats, server in stats_and_servers:
@@ -223,7 +228,7 @@ def assign(
             excluded_servers_uuids: If not None, exclude this servers from selection. Used in case we check the availability of a server
                                  with some external method and we want to exclude it from selection because it has already failed.
             weight_threshold: If not 0, basically will prefer values below an near this value
-            
+
             Note:
                 weight_threshold is used to select a server with a weight as near as possible, without going over, to this value.
                 If none is found, the server with the lowest weight will be selected.
@@ -233,7 +238,7 @@ def assign(
                    * if weight is over threshold, 1 + weight is returned (so, all values over threshold are worse than any value under threshold)
                    that is:
                     real_weight = weight_threshold - weight if weight < weight_threshold else 1 + weight
-                    
+
                 The idea behind this is to be able to select a server not fully empty, but also not fully loaded, so it can be used
                 to leave servers empty as soon as possible, but also to not overload servers that are near to be full.
 

server/src/uds/core/types/deferred_deletion.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+
+#
+# Copyright (c) 2024 Virtual Cable S.L.U.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without modification,
+# are permitted provided that the following conditions are met:
+#
+#    * Redistributions of source code must retain the above copyright notice,
+#      this list of conditions and the following disclaimer.
+#    * Redistributions in binary form must reproduce the above copyright notice,
+#      this list of conditions and the following disclaimer in the documentation
+#      and/or other materials provided with the distribution.
+#    * Neither the name of Virtual Cable S.L.U. nor the names of its contributors
+#      may be used to endorse or promote products derived from this software
+#      without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""
+Author: Adolfo Gómez, dkmaster at dkmon dot com
+"""
+import enum
+
+
+class DeferredStorageGroup(enum.StrEnum):
+    TO_STOP = 'to_stop'
+    STOPPING = 'stopping'
+    TO_DELETE = 'to_delete'
+    DELETING = 'deleting'
+
+    @staticmethod
+    def from_str(value: str) -> 'DeferredStorageGroup':
+        return DeferredStorageGroup(value)

server/src/uds/core/util/storage.py

@@ -29,6 +29,7 @@
 """
 Author: Adolfo Gómez, dkmaster at dkmon dot com
 """
+import contextlib
 import pickle  # nosec: This is e controled pickle use
 import base64
 import hashlib
@@ -198,6 +199,9 @@ def get(self, key: str, default: typing.Any = None) -> typing.Any:
     def delete(self, key: str) -> None:
         self.__delitem__(key)  # pylint: disable=unnecessary-dunder-call
 
+    def clear(self) -> None:
+        self._filtered.delete()  # Removes all keys
+
     # Custom utility methods
     @property
     def group(self) -> str:
@@ -208,39 +212,6 @@ def group(self, value: str) -> None:
         self._group = value or ''
 
 
-class StorageAccess:
-    """
-    Allows the access to the storage as a dict, with atomic transaction if requested
-    """
-
-    _owner: str
-    _group: typing.Optional[str]
-    _atomic: typing.Optional[transaction.Atomic]
-
-    def __init__(
-        self,
-        owner: str,
-        group: typing.Optional[str] = None,
-        atomic: bool = False,
-    ):
-        self._owner = owner
-        self._group = group
-        self._atomic = transaction.atomic() if atomic else None
-
-    def __enter__(self) -> StorageAsDict:
-        if self._atomic:
-            self._atomic.__enter__()
-        return StorageAsDict(
-            owner=self._owner,
-            group=self._group,
-            atomic=bool(self._atomic),
-        )
-
-    def __exit__(self, exc_type: typing.Any, exc_value: typing.Any, traceback: typing.Any) -> None:
-        if self._atomic:
-            self._atomic.__exit__(exc_type, exc_value, traceback)
-
-
 class Storage:
     _owner: str
     _bownwer: bytes
@@ -382,12 +353,17 @@ def remove(
         except Exception:  # nosec: Not interested in processing exceptions, just ignores it
             pass
 
+    @contextlib.contextmanager
     def as_dict(
         self,
         group: typing.Optional[str] = None,
         atomic: bool = False,
-    ) -> StorageAccess:
-        return StorageAccess(self._owner, group=group, atomic=atomic)
+    ) -> typing.Iterator[StorageAsDict]:
+        if atomic:
+            with transaction.atomic():
+                yield StorageAsDict(self._owner, group=group, atomic=True)
+        else:
+            yield StorageAsDict(self._owner, group=group, atomic=False)
 
     def search_by_attr1(
         self, attr1: typing.Union[collections.abc.Iterable[str], str]

server/src/uds/models/user.py

@@ -210,10 +210,14 @@ def pre_delete(sender: typing.Any, **kwargs: typing.Any) -> None:  # pylint: dis
         to_delete.clean_related_data()
 
         # Remove related stored values
-        with storage.StorageAccess('manager' + str(to_delete.manager.uuid)) as store:
-            for key in store.keys():
-                store.delete(key)
-
+        try:
+            storage.StorageAsDict(
+                owner='manager' + str(to_delete.manager.uuid),
+                group=None,
+                atomic=False
+            ).clear()
+        except Exception:
+            logger.exception('Removing stored data')
         # now removes all "child" of this user, if it has children
         User.objects.filter(parent=to_delete.id).delete()