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

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.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: "AnyGateway" 

41 log: logging.Logger 

42 

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

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

45 self._updating_info = False 

46 self.__deserialize() 

47 

48 def __deserialize(self) -> None: 

49 if self.stored.extra_attributes is not None: 

50 self.deserialize_extra_attributes(self.stored.extra_attributes) 

51 

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

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

54 orm.add(self.stored) 

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

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

57 self.__deserialize() 

58 

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

60 """ 

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

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

63 `deserialize_extra_attributes()`. 

64 

65 """ 

66 return None 

67 

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

69 """ 

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

71 `serialize_extra_attributes()`. 

72 

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

74 Be sure to convert to other types if necessary. 

75 """ 

76 pass 

77 

78 @contextmanager 

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

80 self._updating_info = True 

81 yield 

82 self._updating_info = False 

83 self.stored.updated = True 

84 self.commit() 

85 

86 def commit(self) -> None: 

87 if self._updating_info: 

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

89 else: 

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

91 super().commit() 

92 

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

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

95 setattr(self.stored, key, value) 

96 if self._updating_info: 

97 return 

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

99 orm.execute( 

100 sa.update(self.stored.__class__) 

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

102 .values(**kwargs) 

103 ) 

104 orm.commit()