Coverage for slidge/db/meta.py: 77%

1from __future__ import annotations 

2 

3import json 

4from collections.abc import Mapping 

5from typing import Any, ClassVar 

6 

7import sqlalchemy as sa 

8import sqlalchemy.orm 

9from slixmpp import JID 

10from sqlalchemy import Dialect 

11 

12 

13class JIDType(sa.TypeDecorator[JID]): 

14 """ 

15 Custom SQLAlchemy type for JIDs 

16 """ 

17 

18 impl = sa.types.TEXT 

19 cache_ok = True 

20 

21 def process_bind_param(self, value: JID | None, dialect: sa.Dialect) -> str | None: 

22 if value is None: 

23 return value 

24 return str(value) 

25 

26 def process_result_value( 

27 self, value: str | None, dialect: sa.Dialect 

28 ) -> JID | None: 

29 if value is None: 

30 return value 

31 return JID(value) 

32 

33 

34type JSONSerializableTypes = ( 

35 str | float | bool | Mapping[str, "JSONSerializableTypes"] | None 

36) 

37JSONSerializable = dict[str | int, JSONSerializableTypes] 

38 

39 

40_INT_KEY_TAG = "\x00" 

41 

42 

43def _encode_key(key: str | int) -> str: 

44 if isinstance(key, int): 

45 return f"{_INT_KEY_TAG}{key}" 

46 return key 

47 

48 

49def _decode_key(key: str) -> str | int: 

50 if key.startswith(_INT_KEY_TAG): 

51 return int(key[1:]) 

52 return key 

53 

54 

55def _encode(obj: Any) -> Any: # noqa: ANN401 

56 if isinstance(obj, dict): 

57 return {_encode_key(k): _encode(v) for k, v in obj.items()} 

58 if isinstance(obj, list): 

59 return [_encode(item) for item in obj] 

60 return obj 

61 

62 

63def _decode(obj: Any) -> Any: # noqa: ANN401 

64 if isinstance(obj, dict): 

65 return {_decode_key(k): _decode(v) for k, v in obj.items()} 

66 if isinstance(obj, list): 

67 return [_decode(item) for item in obj] 

68 return obj 

69 

70 

71class JSONEncodedDict(sa.TypeDecorator[JSONSerializable]): 

72 """ 

73 Custom SQLAlchemy type for dictionaries stored as JSON 

74 

75 Note that mutations of the dictionary are not detected by SQLAlchemy, 

76 which is why use ``attributes.flag_modified()`` in ``UserStore.update()`` 

77 

78 Integer keys for dicts are supported. Internally, they are prefixed will 

79 null bytes to mark them as integers. As a consequence, null byte as the 

80 first character of a string key is NOT SUPPORTED. 

81 """ 

82 

83 impl = sa.VARCHAR 

84 

85 cache_ok = True 

86 

87 def process_bind_param( 

88 self, value: JSONSerializable | None, dialect: Dialect 

89 ) -> str | None: 

90 if value is None: 

91 return None 

92 return json.dumps(_encode(value)) 

93 

94 def process_result_value( 

95 self, value: object, dialect: Dialect 

96 ) -> JSONSerializable | None: 

97 if value is None: 

98 return None 

99 return _decode(json.loads(value)) # type:ignore[no-any-return,arg-type] 

100 

101 

102class Base(sa.orm.DeclarativeBase): 

103 type_annotation_map: ClassVar[dict[object, object]] = { 

104 JSONSerializable: JSONEncodedDict, 

105 JID: JIDType, 

106 } 

107 

108 

109Base.metadata.naming_convention = { 

110 "ix": "ix_%(column_0_label)s", 

111 "uq": "uq_%(table_name)s_%(column_0_name)s", 

112 "ck": "ck_%(table_name)s_`%(constraint_name)s`", 

113 "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", 

114 "pk": "pk_%(table_name)s", 

115} 

116 

117 

118def get_engine(path: str, echo: bool = False, pool_size: int = 5) -> sa.Engine: 

119 from sqlalchemy import log as sqlalchemy_log 

120 

121 engine = sa.create_engine( 

122 path, 

123 pool_size=pool_size, 

124 json_serializer=lambda obj: json.dumps(obj, ensure_ascii=False), 

125 ) 

126 if echo: 

127 sqlalchemy_log._add_default_handler = lambda x: None # type:ignore 

128 engine.echo = True 

129 return engine 

130 

131 

132def modified_attributes(instance: Base) -> set[str]: 

133 attrs = set() 

134 state = sa.inspect(instance) 

135 for attribute in state.attrs: 

136 if attribute.key in ("avatar", "banner"): 

137 # avatar and banner are pubsub *items*, not part of the 

138 # pubsub#metadata form; broadcasting their update has a 

139 # different protocol 

140 continue 

141 history = attribute.history 

142 if history.has_changes(): 

143 attrs.add(attribute.key) 

144 return attrs