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

1import logging 

2import typing 

3from contextlib import contextmanager 

4 

5import sqlalchemy as sa 

6 

7from ...db.meta import Base, JSONSerializable 

8from ...db.models import Contact, Room 

9 

10if typing.TYPE_CHECKING: 

11 from slidge.util.types import AnyGateway 

12 

13 

14class DBMixin: 

15 stored: Base 

16 xmpp: "AnyGateway" 

17 log: logging.Logger 

18 

19 def merge(self) -> None: 

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

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

22 

23 def commit(self) -> None: 

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

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

26 orm.add(self.stored) 

27 orm.commit() 

28 

29 

30class UpdateInfoMixin(DBMixin): 

31 """ 

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

33 on every attribute change. 

34 """ 

35 

36 stored: Contact | Room 

37 xmpp: "AnyGateway" 

38 log: logging.Logger 

39 

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

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

42 self._updating_info = False 

43 self.__deserialize() 

44 

45 def __deserialize(self) -> None: 

46 if self.stored.extra_attributes is not None: 

47 self.deserialize_extra_attributes(self.stored.extra_attributes) 

48 

49 def refresh(self, attrs: list[str] | None = None) -> None: 

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

51 orm.add(self.stored) 

52 orm.refresh(self.stored, attribute_names=attrs) 

53 if attrs is None or "extra_attributes" in attrs: 

54 self.__deserialize() 

55 

56 def serialize_extra_attributes(self) -> JSONSerializable | None: 

57 """ 

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

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

60 `deserialize_extra_attributes()`. 

61 

62 """ 

63 return None 

64 

65 def deserialize_extra_attributes(self, data: JSONSerializable) -> None: 

66 """ 

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

68 `serialize_extra_attributes()`. 

69 

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

71 Be sure to convert to other types if necessary. 

72 """ 

73 

74 @contextmanager 

75 def updating_info(self) -> typing.Iterator[None]: 

76 self._updating_info = True 

77 yield 

78 self._updating_info = False 

79 self.stored.updated = True 

80 self.commit() 

81 

82 def commit(self) -> None: 

83 if self._updating_info: 

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

85 else: 

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

87 super().commit() 

88 

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

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

91 setattr(self.stored, key, value) 

92 if self._updating_info: 

93 return 

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

95 orm.execute( 

96 sa.update(self.stored.__class__) 

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

98 .values(**kwargs) 

99 ) 

100 orm.commit()