Coverage for slidge/core/mixins/db.py: 94%

1import logging 

2import typing 

3from contextlib import AbstractContextManager as ContextManager 

4from contextlib import contextmanager, nullcontext 

5 

6import sqlalchemy as sa 

7import sqlalchemy.orm 

8 

9from ...db.models import Contact, Participant, Room, Space 

10from ...db.track import TrackedDict 

11from ...util.types import ExtraAttrMapping 

12from .base import SessionBound 

13 

14if typing.TYPE_CHECKING: 

15 from slidge.util.types import AnyGateway 

16 

17 

18class UpdateInfoMixin[ 

19 T: ExtraAttrMapping = typing.Any, 

20 PropsT = typing.Mapping[str, typing.Any], 

21](SessionBound): 

22 _stored: Contact | Room | Participant | Space 

23 xmpp: "AnyGateway" 

24 log: logging.Logger 

25 _batch_set_attributes_enabled = False 

26 

27 def __init__(self, *args: object, **kwargs: object) -> None: 

28 super().__init__(*args, **kwargs) 

29 self._pending_updates: dict[str, object] = {} 

30 self._extra_attributes = typing.cast( 

31 T, 

32 TrackedDict( 

33 typing.cast(ExtraAttrMapping, self._stored.extra_attributes or {}), 

34 self._on_extra_attr_change, 

35 ), 

36 ) 

37 

38 @classmethod 

39 @contextmanager 

40 def _batch_set_attributes(cls) -> typing.Generator[None]: 

41 cls._batch_set_attributes_enabled = True 

42 try: 

43 yield 

44 finally: 

45 cls._batch_set_attributes_enabled = False 

46 

47 def batch_set_attributes( 

48 self, 

49 **kwargs: typing.Unpack[PropsT], # type:ignore[misc] 

50 ) -> ContextManager[None]: 

51 if not kwargs: 

52 return self.__batch_set_attributes_context() 

53 

54 with self.__batch_set_attributes_context(): 

55 for name, value in kwargs.items(): 

56 setattr(self, name, value) 

57 return nullcontext() 

58 

59 @contextmanager 

60 def __batch_set_attributes_context(self) -> typing.Generator[None]: 

61 self._batch_set_attributes_enabled = True 

62 try: 

63 yield 

64 finally: 

65 self._batch_set_attributes_enabled = False 

66 if self._pending_updates: 

67 with self.xmpp.store.session() as orm: 

68 orm.execute( 

69 sa.update(self._stored.__class__) 

70 .where(self._stored.__class__.id == self._stored.id) 

71 .values(**self._pending_updates) 

72 ) 

73 orm.commit() 

74 

75 def _update_stored_attribute(self, **kwargs: object) -> bool: 

76 kwargs = self._stored._normalize_stored_updates(kwargs) 

77 for key, value in list(kwargs.items()): 

78 try: 

79 change = getattr(self._stored, key) != value 

80 except sa.orm.exc.DetachedInstanceError: 

81 change = True 

82 if change: 

83 self._stored._set_committed_value(key, value) 

84 else: 

85 kwargs.pop(key) 

86 if self._batch_set_attributes_enabled: 

87 self._pending_updates.update(kwargs) 

88 return False 

89 if kwargs: 

90 with self.xmpp.store.session() as orm: 

91 orm.execute( 

92 sa.update(self._stored.__class__) 

93 .where(self._stored.__class__.id == self._stored.id) 

94 .values(**kwargs) 

95 ) 

96 orm.commit() 

97 return True 

98 else: 

99 return False 

100 

101 @property 

102 def extra_attributes(self) -> T: 

103 return self._extra_attributes 

104 

105 def _on_extra_attr_change(self, new_data: T) -> None: 

106 self._update_stored_attribute(extra_attributes=new_data)