Coverage for slidge/db/alembic/versions/9c5a127d4600_attachments_in_several_columns.py: 69%

1"""attachments in several columns 

2 

3Revision ID: 9c5a127d4600 

4Revises: 7d21996a1605 

5Create Date: 2026-09-20 20:41:30.934062 

6 

7""" 

8 

9import logging 

10from collections.abc import Sequence 

11from xml.etree import ElementTree as ET 

12 

13import slixmpp.plugins.xep_0446.stanza 

14import slixmpp.plugins.xep_0447.stanza 

15import sqlalchemy as sa 

16from alembic import op 

17from slixmpp import Message, register_stanza_plugin 

18 

19# revision identifiers, used by Alembic. 

20revision: str = "9c5a127d4600" 

21down_revision: str | None = "7d21996a1605" 

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

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

24 

25log = logging.getLogger("Attachment migration") 

26 

27 

28def upgrade() -> None: 

29 _register_stanza_plugins() 

30 

31 with op.batch_alter_table("attachment", schema=None) as batch_op: 

32 batch_op.add_column(sa.Column("size", sa.Integer(), nullable=True)) 

33 batch_op.add_column(sa.Column("sha256", sa.String(), nullable=True)) 

34 batch_op.add_column(sa.Column("thumbhash", sa.String(), nullable=True)) 

35 batch_op.add_column(sa.Column("width", sa.Integer(), nullable=True)) 

36 batch_op.add_column(sa.Column("height", sa.Integer(), nullable=True)) 

37 batch_op.add_column(sa.Column("content_type", sa.String(), nullable=True)) 

38 batch_op.add_column(sa.Column("desc", sa.String(), nullable=True)) 

39 batch_op.add_column(sa.Column("name", sa.String(), nullable=True)) 

40 batch_op.add_column(sa.Column("date", sa.DateTime(), nullable=True)) 

41 batch_op.add_column( 

42 sa.Column( 

43 "disposition", 

44 sa.Enum("inline", "attachment", native_enum=False), 

45 nullable=True, 

46 ) 

47 ) 

48 

49 conn = op.get_bind() 

50 q = sa.select(att_tbl.c.id, att_tbl.c.sfs) 

51 update_stmt = att_tbl.update().where(att_tbl.c.id == sa.bindparam("_id")) 

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

53 

54 to_delete: list[int] = [] 

55 for batch in result.partitions(): 

56 params = [] 

57 for row in batch: 

58 id_, sfs_str = row 

59 try: 

60 sfs = slixmpp.plugins.xep_0447.stanza.StatelessFileSharing( 

61 xml=ET.fromstring(sfs_str) 

62 ) 

63 sha256 = sfs["file"]["hash"]["value"] if "hash" in sfs["file"] else None 

64 if "thumbnail" in sfs["file"]: 

65 thumbhash = sfs["file"]["thumbnail"]["uri"].removeprefix( 

66 "data:image/thumbhash;base64," 

67 ) 

68 width = sfs["file"]["thumbnail"]["width"] or None 

69 height = sfs["file"]["thumbnail"]["height"] or None 

70 else: 

71 thumbhash = width = height = None 

72 params.append( 

73 { 

74 "_id": id_, 

75 "size": sfs["file"]["size"] or None, 

76 "width": width, 

77 "height": height, 

78 "content_type": sfs["file"]["media-type"] or None, 

79 "desc": sfs["file"]["desc"] or None, 

80 "name": sfs["file"]["name"] or None, 

81 "date": sfs["file"]["date"] or None, 

82 "sha256": sha256, 

83 "thumbhash": thumbhash, 

84 "disposition": sfs["disposition"] or None, 

85 } 

86 ) 

87 except Exception: 

88 log.exception("Could not migrate row: %s", sfs_str) 

89 to_delete.append(id_) 

90 continue 

91 

92 if params: 

93 conn.execute(update_stmt, params) 

94 

95 if to_delete: 

96 conn.execute(att_tbl.delete().where(att_tbl.c.id.in_(to_delete))) 

97 

98 with op.batch_alter_table("attachment", schema=None) as batch_op: 

99 batch_op.drop_column("sfs") 

100 batch_op.drop_column("sims") 

101 

102 

103def _register_stanza_plugins() -> None: 

104 slixmpp.plugins.xep_0446.stanza.register_plugins() 

105 register_stanza_plugin( 

106 Message, slixmpp.plugins.xep_0447.stanza.StatelessFileSharing 

107 ) 

108 

109 register_stanza_plugin( 

110 slixmpp.plugins.xep_0447.stanza.StatelessFileSharing, 

111 slixmpp.plugins.xep_0447.stanza.Sources, 

112 ) 

113 register_stanza_plugin( 

114 slixmpp.plugins.xep_0447.stanza.StatelessFileSharing, 

115 slixmpp.plugins.xep_0446.stanza.File, 

116 ) 

117 register_stanza_plugin( 

118 slixmpp.plugins.xep_0447.stanza.Sources, 

119 slixmpp.plugins.xep_0447.stanza.UrlData, 

120 iterable=True, 

121 ) 

122 

123 

124def downgrade() -> None: 

125 raise RuntimeError("downgrades are not supported") 

126 

127 

128att_tbl = sa.table( 

129 "attachment", 

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

131 sa.column("sfs", sa.String), 

132 sa.Column("size", sa.Integer(), nullable=True), 

133 sa.Column("sha256", sa.String(), nullable=True), 

134 sa.Column("thumbhash", sa.String(), nullable=True), 

135 sa.Column("width", sa.Integer(), nullable=True), 

136 sa.Column("height", sa.Integer(), nullable=True), 

137 sa.Column("content_type", sa.String(), nullable=True), 

138 sa.Column("desc", sa.String(), nullable=True), 

139 sa.Column("name", sa.String(), nullable=True), 

140 sa.Column("disposition", sa.String(), nullable=True), 

141 sa.Column("date", sa.DateTime(), nullable=True), 

142)