mirror of
https://github.com/dkmstr/openuds.git
synced 2025-03-30 18:50:20 +03:00
Added dict-like storage management
This commit is contained in:
Adolfo Gómez García
parent
dd39bb4e64
commit
37f06617b8
@ -37,12 +37,14 @@ import base64
|
||||
import hashlib
|
||||
from collections.abc import MutableMapping
|
||||
|
||||
from django.db import transaction
|
||||
from django.db import transaction, models
|
||||
from uds.models.storage import Storage as DBStorage
|
||||
from uds.core.util import encoders
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MARK = '_mgb_'
|
||||
|
||||
def _calcKey(owner: bytes, key: bytes, extra: typing.Optional[bytes] = None) -> str:
|
||||
h = hashlib.md5()
|
||||
h.update(owner)
|
||||
@ -51,16 +53,19 @@ def _calcKey(owner: bytes, key: bytes, extra: typing.Optional[bytes] = None) ->
|
||||
h.update(extra)
|
||||
return h.hexdigest()
|
||||
|
||||
def _encodeValue(key: str, value: typing.Any) -> str:
|
||||
return base64.b64encode(pickle.dumps((key, value))).decode()
|
||||
def _encodeValue(key: str, value: typing.Any, compat: bool = False) -> str:
|
||||
if not compat:
|
||||
return base64.b64encode(pickle.dumps((MARK, key, value))).decode()
|
||||
# Compatibility save
|
||||
return base64.b64encode(pickle.dumps(value)).decode()
|
||||
|
||||
def _decodeValue(dbk: str, value: typing.Optional[str]) -> typing.Tuple[str, typing.Any]:
|
||||
if value:
|
||||
try:
|
||||
v = pickle.loads(base64.b64decode(value.encode()))
|
||||
if isinstance(v, tuple):
|
||||
return typing.cast(typing.Tuple[str, typing.Any], v)
|
||||
# Fix value so it contains also the "key" (in this case, no valid key...)
|
||||
if isinstance(v, tuple) and v[0] == MARK:
|
||||
return typing.cast(typing.Tuple[str, typing.Any], v[1:])
|
||||
# Fix value so it contains also the "key" (in this case, the original key is lost, we have only the hash value...)
|
||||
return ('#' + dbk, v)
|
||||
except Exception as e:
|
||||
logger.warn('Unknown pickable value: %s (%s)', value, e)
|
||||
@ -71,16 +76,39 @@ class StorageAsDict(MutableMapping):
|
||||
'''
|
||||
Accesses storage as dictionary. Much more convenient that old method
|
||||
'''
|
||||
def __init__(self, owner: str, group: typing.Optional[str], atomic: bool = False) -> None:
|
||||
def __init__(self, owner: str, group: typing.Optional[str], atomic: bool = False, compat: bool = False) -> None:
|
||||
"""Initializes an storage as dict accesor
|
||||
|
||||
Args:
|
||||
owner (str): owner of the storage
|
||||
group (typing.Optional[str]): group for this dict
|
||||
atomic (bool, optional): [description]. if True, operations with DB will be atomic
|
||||
compat (bool, optional): [description]. if True, keys will be generated "old way" (ignoring group)
|
||||
"""
|
||||
self._group = group or ''
|
||||
self._owner = owner
|
||||
self._atomic = atomic # Not used right now, maybe removed
|
||||
self._compat = compat
|
||||
|
||||
@property
|
||||
def _db(self) -> typing.Union[models.QuerySet, models.Manager]:
|
||||
if self._atomic:
|
||||
return DBStorage.objects.select_for_update()
|
||||
else:
|
||||
return DBStorage.objects
|
||||
|
||||
@property
|
||||
def _filtered(self) -> models.QuerySet:
|
||||
fltr = self._db.filter(owner=self._owner)
|
||||
if self._group:
|
||||
fltr = fltr.filter(attr1=self._group)
|
||||
return fltr
|
||||
|
||||
def _key(self, key: str) -> str:
|
||||
if key[0] == '#':
|
||||
# Compat with old dbk
|
||||
# Compat with old db key
|
||||
return key[1:]
|
||||
return _calcKey(self._owner.encode(), key.encode(), self._group.encode())
|
||||
return _calcKey(self._owner.encode(), key.encode())
|
||||
|
||||
def __getitem__(self, key: str) -> typing.Any:
|
||||
if not isinstance(key, str):
|
||||
@ -89,7 +117,7 @@ class StorageAsDict(MutableMapping):
|
||||
dbk = self._key(key)
|
||||
logger.debug('Getitem: %s', dbk)
|
||||
try:
|
||||
c: DBStorage = DBStorage.objects.get(pk=dbk)
|
||||
c: DBStorage = self._db.get(pk=dbk)
|
||||
return _decodeValue(dbk, c.data)[1] # Ignores original key
|
||||
except DBStorage.DoesNotExist:
|
||||
return None
|
||||
@ -100,44 +128,59 @@ class StorageAsDict(MutableMapping):
|
||||
|
||||
dbk = self._key(key)
|
||||
logger.debug('Setitem: %s = %s', dbk, value)
|
||||
data = _encodeValue(key, value)
|
||||
data = _encodeValue(key, value, self._compat)
|
||||
c, created = DBStorage.objects.update_or_create(key=dbk, defaults={'data': data, 'attr1': self._group, 'owner': self._owner})
|
||||
|
||||
def __delitem__(self, key: str):
|
||||
def __delitem__(self, key: str) -> None:
|
||||
dbk = self._key(key)
|
||||
logger.debug('Delitem: %s', key)
|
||||
logger.debug('Delitem: %s --> %s', key, dbk)
|
||||
DBStorage.objects.filter(key=dbk).delete()
|
||||
|
||||
def __iter__(self):
|
||||
'''
|
||||
Iterates through keys
|
||||
'''
|
||||
return iter((_decodeValue(i.key, i.data)[0] for i in DBStorage.objects.filter(owner=self._owner, attr1=self._group)))
|
||||
return iter(_decodeValue(i.key, i.data)[0] for i in self._filtered)
|
||||
|
||||
def __contains__(self, key: object) -> bool:
|
||||
logger.debug('Contains: %s', key)
|
||||
if isinstance(key, str):
|
||||
dbk = self._key(key)
|
||||
return DBStorage.objects.filter(owner=self._owner, attr1=self._group, key=dbk).count() > 0
|
||||
return self._filtered.filter(key=self._key(key)).exists()
|
||||
return False
|
||||
|
||||
def __len__(self):
|
||||
return DBStorage.objects.filter(owner=self._owner, attr1=self._group).count()
|
||||
return self._filtered.count()
|
||||
|
||||
# Optimized methods, avoid re-reading from DB
|
||||
def items(self):
|
||||
return iter(_decodeValue(i.key, i.data) for i in self._filtered)
|
||||
|
||||
def values(self):
|
||||
return iter(_decodeValue(i.key, i.data)[1] for i in self._filtered)
|
||||
|
||||
# Custom utility methods
|
||||
@property
|
||||
def group(self) -> str:
|
||||
return self._group or ''
|
||||
|
||||
@group.setter
|
||||
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
|
||||
'''
|
||||
def __init__(self, owner: str, group: typing.Optional[str] = None, atomic: typing.Optional[bool] = False):
|
||||
def __init__(self, owner: str, group: typing.Optional[str] = None, atomic: bool = False, compat: bool = False):
|
||||
self._owner = owner
|
||||
self._group = group
|
||||
self._atomic = transaction.atomic() if atomic else None
|
||||
self._compat = compat
|
||||
|
||||
def __enter__(self):
|
||||
if self._atomic:
|
||||
self._atomic.__enter__()
|
||||
return StorageAsDict(self._owner, self._group, bool(self._atomic))
|
||||
return StorageAsDict(owner=self._owner, group=self._group, atomic=bool(self._atomic), compat=self._compat)
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
if self._atomic:
|
||||
@ -237,8 +280,8 @@ class Storage:
|
||||
"""
|
||||
# dbStorage.objects.unlock() # @UndefinedVariable
|
||||
|
||||
def map(self, group: typing.Optional[str] = None, atomic: typing.Optional[bool] = False):
|
||||
return StorageAccess(self._owner, group, atomic)
|
||||
def map(self, group: typing.Optional[str] = None, atomic: bool = False, compat: bool = False) -> StorageAccess:
|
||||
return StorageAccess(self._owner, group=group, atomic=atomic, compat=compat)
|
||||
|
||||
def locateByAttr1(self, attr1: typing.Union[typing.Iterable[str], str]) -> typing.Iterable[bytes]:
|
||||
if isinstance(attr1, str):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user