Coverage for slidge/db/alembic/versions/185cc8356e10_mam_do_not_store_full_stanza.py: 61%

1"""mam: do not store full stanza 

2 

3Revision ID: 185cc8356e10 

4Revises: c594846b3e03 

5Create Date: 2026-09-08 23:01:56.978966 

6 

7""" 

8 

9import logging 

10import xml.etree.ElementTree as ET 

11from collections.abc import Sequence 

12 

13import slixmpp.plugins.xep_0045.stanza 

14import slixmpp.plugins.xep_0203.stanza 

15import slixmpp.plugins.xep_0359.stanza 

16import slixmpp.plugins.xep_0421.stanza 

17import sqlalchemy as sa 

18from alembic import op 

19from slixmpp import Message, register_stanza_plugin 

20from slixmpp.plugins.xep_0333.stanza import Displayed 

21from slixmpp.plugins.xep_0334 import NoPermanentStore, NoStore, Store 

22from slixmpp.plugins.xep_0424.stanza import Retract 

23 

24from slidge.db.store import _serialize_children 

25 

26# revision identifiers, used by Alembic. 

27revision: str = "185cc8356e10" 

28down_revision: str | None = "c594846b3e03" 

29branch_labels: str | Sequence[str] | None = None 

30depends_on: str | Sequence[str] | None = None 

31 

32 

33def upgrade() -> None: 

34 with op.batch_alter_table("mam", schema=None) as batch_op: 

35 batch_op.add_column(sa.Column("author_affiliation", sa.String(), nullable=True)) 

36 batch_op.add_column(sa.Column("author_role", sa.String(), nullable=True)) 

37 batch_op.add_column( 

38 sa.Column("author_jid_localpart", sa.String(), nullable=True) 

39 ) 

40 batch_op.add_column(sa.Column("occupant_id", sa.String(), nullable=True)) 

41 batch_op.add_column(sa.Column("payload", sa.String(), nullable=True)) 

42 batch_op.add_column(sa.Column("from_user", sa.Boolean(), nullable=True)) 

43 

44 conn = op.get_bind() 

45 

46 q = sa.select( 

47 mam_tbl.c.id, 

48 mam_tbl.c.stanza, 

49 user_tbl.c.jid, 

50 ).select_from( 

51 mam_tbl.join(room_tbl, mam_tbl.c.room_id == room_tbl.c.id).join( 

52 user_tbl, room_tbl.c.user_account_id == user_tbl.c.id 

53 ) 

54 ) 

55 

56 update_stmt = mam_tbl.update().where(mam_tbl.c.id == sa.bindparam("_id")) 

57 

58 count_q = sa.select(sa.func.count()).select_from(q.subquery()) 

59 total = conn.execute(count_q).scalar_one() 

60 log.info("Migrating %s rows...", total) 

61 

62 result = conn.execute(q, execution_options={"yield_per": 1000}) 

63 for i, batch in enumerate(result.partitions()): 

64 log.debug("Batch %s", i) 

65 params = [] 

66 for row in batch: 

67 mam_id, stanza, user_jid = row 

68 try: 

69 values = _parse(stanza, user_jid) 

70 except Exception: 

71 log.exception("Failed parsing row %s", mam_id) 

72 values = { 

73 "author_affiliation": "member", 

74 "author_role": "participant", 

75 "author_jid_localpart": "failed-migration-jid", 

76 "occupant_id": "failed-migration-occupant", 

77 "from_user": False, 

78 "payload": "", 

79 } 

80 values["_id"] = mam_id 

81 params.append(values) 

82 

83 if params: 

84 conn.execute(update_stmt, params) 

85 

86 with op.batch_alter_table("mam", schema=None) as batch_op: 

87 batch_op.alter_column("from_user", nullable=False) 

88 batch_op.alter_column("payload", nullable=False) 

89 batch_op.drop_column("stanza") 

90 

91 

92def downgrade() -> None: 

93 raise RuntimeError("downgrades are not supported") 

94 

95 

96def _parse(stanza_str: str, user_jid: str) -> dict[str, str | bool | None]: 

97 stanza = Message(xml=ET.fromstring(stanza_str)) 

98 

99 author_affiliation = stanza["muc"]["affiliation"] or None 

100 author_role = stanza["muc"]["role"] or None 

101 author_jid = stanza["muc"]["jid"] or None 

102 

103 if author_jid == user_jid: 

104 from_user = True 

105 author_jid_localpart = None 

106 else: 

107 from_user = False 

108 author_jid_localpart = author_jid.user if author_jid else None 

109 

110 occupant_id = stanza["occupant-id"]["id"] or None 

111 

112 del stanza["delay"] 

113 del stanza["markable"] 

114 del stanza["store"] 

115 del stanza["chat_state"] 

116 del stanza["muc"] 

117 del stanza["origin_id"] 

118 del stanza["stanza_id"] 

119 del stanza["occupant-id"] 

120 

121 payload = _serialize_children(stanza) 

122 

123 return { 

124 "author_affiliation": author_affiliation, 

125 "author_role": author_role, 

126 "author_jid_localpart": author_jid_localpart, 

127 "occupant_id": occupant_id, 

128 "from_user": from_user, 

129 "payload": payload, 

130 } 

131 

132 

133mam_tbl = sa.table( 

134 "mam", 

135 sa.column("id", sa.Integer), 

136 sa.column("room_id", sa.Integer), 

137 sa.column("stanza", sa.String), 

138 sa.column("payload", sa.String), 

139 sa.column("author_affiliation", sa.String), 

140 sa.column("author_role", sa.String), 

141 sa.column("author_jid_localpart", sa.String), 

142 sa.column("occupant_id", sa.String), 

143 sa.column("from_user", sa.Boolean), 

144) 

145 

146room_tbl = sa.table( 

147 "room", 

148 sa.column("id", sa.Integer), 

149 sa.column("user_account_id", sa.Integer), 

150) 

151user_tbl = sa.table( 

152 "user_account", 

153 sa.column("id", sa.Integer), 

154 sa.column("jid", sa.String), 

155) 

156 

157register_stanza_plugin(Message, Displayed) 

158register_stanza_plugin(Message, Retract) 

159register_stanza_plugin(Message, Store) 

160register_stanza_plugin(Message, NoStore) 

161register_stanza_plugin(Message, NoPermanentStore) 

162slixmpp.plugins.xep_0045.stanza.register_plugins() 

163register_stanza_plugin(Message, slixmpp.plugins.xep_0203.stanza.Delay) 

164slixmpp.plugins.xep_0359.stanza.register_plugins() 

165register_stanza_plugin(Message, slixmpp.plugins.xep_0421.stanza.OccupantId) 

166 

167log = logging.getLogger("MAM migration")