Coverage for mindsdb / utilities / context.py: 79%

1from contextvars import ContextVar 

2from typing import Any 

3from copy import deepcopy 

4 

5 

6class Context: 

7 """Thread independent storage""" 

8 

9 __slots__ = ("_storage",) 

10 

11 def __init__(self, storage) -> None: 

12 object.__setattr__(self, "_storage", storage) 

13 self.set_default() 

14 

15 def set_default(self) -> None: 

16 self._storage.set( 

17 { 

18 "user_id": None, 

19 "company_id": None, 

20 "session_id": "", 

21 "task_id": None, 

22 "user_class": 0, 

23 "profiling": {"level": 0, "enabled": False, "pointer": None, "tree": None}, 

24 "email_confirmed": 0, 

25 "used_handlers": set(), 

26 "params": {}, 

27 } 

28 ) 

29 

30 def __getattr__(self, name: str) -> Any: 

31 storage = self._storage.get({}) 

32 if name not in storage: 

33 raise AttributeError(name) 

34 return storage[name] 

35 

36 def __setattr__(self, name: str, value: Any) -> None: 

37 storage = deepcopy(self._storage.get({})) 

38 storage[name] = value 

39 self._storage.set(storage) 

40 

41 def __delattr__(self, name: str) -> None: 

42 storage = deepcopy(self._storage.get({})) 

43 if name not in storage: 

44 raise AttributeError(name) 

45 del storage["name"] 

46 self._storage.set(storage) 

47 

48 def dump(self) -> dict: 

49 storage = deepcopy(self._storage.get({})) 

50 return storage 

51 

52 def load(self, storage: dict) -> None: 

53 self._storage.set(storage) 

54 

55 def get_metadata(self, **kwargs) -> dict: 

56 return { 

57 "user_id": self.user_id or "", 

58 "company_id": self.company_id or "", 

59 "session_id": self.session_id, 

60 "user_class": self.user_class, 

61 **kwargs, 

62 } 

63 

64 

65_context_var = ContextVar("mindsdb.context") 

66context = Context(_context_var)