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

1import logging 

2import typing 

3from contextlib import contextmanager 

4 

5import sqlalchemy as sa 

6 

7from ...db.models import Base, Contact, Room 

8 

9if typing.TYPE_CHECKING: 

10 from slidge import BaseGateway 

11 

12 

13class DBMixin: 

14 stored: Base 

15 xmpp: "BaseGateway" 

16 log: logging.Logger 

17 

18 def merge(self) -> None: 

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

20 self.stored = orm.merge(self.stored) 

21 

22 def commit(self, merge: bool = False) -> None: 

23 with self.xmpp.store.session(expire_on_commit=False) as orm: 

24 if merge: 

25 self.log.debug("Merging %s", self.stored) 

26 self.stored = orm.merge(self.stored) 

27 self.log.debug("Merged %s", self.stored) 

28 orm.add(self.stored) 

29 self.log.debug("Committing to DB") 

30 orm.commit() 

31 

32 

33class UpdateInfoMixin(DBMixin): 

34 """ 

35 This mixin just adds a context manager that prevents commiting to the DB 

36 on every attribute change. 

37 """ 

38 

39 stored: Contact | Room 

40 xmpp: "BaseGateway" 

41 log: logging.Logger 

42 

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

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

45 self._updating_info = False 

46 self.__deserialize() 

47 

48 def __deserialize(self): 

49 if self.stored.extra_attributes is not None: 

50 self.deserialize_extra_attributes(self.stored.extra_attributes) 

51 

52 def refresh(self) -> None: 

53 with self.xmpp.store.session(expire_on_commit=False) as orm: 

54 orm.add(self.stored) 

55 orm.refresh(self.stored) 

56 self.__deserialize() 

57 

58 def serialize_extra_attributes(self) -> dict | None: 

59 """ 

60 If you want custom attributes of your instance to be stored persistently 

61 to the DB, here is where you have to return them as a dict to be used in 

62 `deserialize_extra_attributes()`. 

63 

64 """ 

65 return None 

66 

67 def deserialize_extra_attributes(self, data: dict) -> None: 

68 """ 

69 This is where you get the dict that you passed in 

70 `serialize_extra_attributes()`. 

71 

72 ⚠ Since it is serialized as json, dictionary keys are converted to strings! 

73 Be sure to convert to other types if necessary. 

74 """ 

75 pass 

76 

77 @contextmanager 

78 def updating_info(self): 

79 self._updating_info = True 

80 yield 

81 self._updating_info = False 

82 self.stored.updated = True 

83 self.commit() 

84 

85 def commit(self, merge: bool = False) -> None: 

86 if self._updating_info: 

87 self.log.debug("Not updating %s right now", self.stored) 

88 else: 

89 self.stored.extra_attributes = self.serialize_extra_attributes() 

90 super().commit(merge=merge) 

91 

92 def update_stored_attribute(self, **kwargs) -> None: 

93 for key, value in kwargs.items(): 

94 setattr(self.stored, key, value) 

95 if self._updating_info: 

96 return 

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

98 orm.execute( 

99 sa.update(self.stored.__class__) 

100 .where(self.stored.__class__.id == self.stored.id) 

101 .values(**kwargs) 

102 ) 

103 orm.commit()