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

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 

77 @contextmanager 

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

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) -> 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() 

91 

92 def update_stored_attribute(self, **kwargs: object) -> 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()