Coverage for slidge/db/store.py: 88%

408 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-09-10 04:45 +0000

1from __future__ import annotations 

2 

3import hashlib 

4import logging 

5import shutil 

6import uuid 

7from collections.abc import Callable, Collection, Iterable, Iterator 

8from datetime import UTC, datetime, timedelta 

9from mimetypes import guess_extension 

10from typing import Any, ClassVar, Protocol, TypeVar 

11 

12import sqlalchemy as sa 

13import sqlalchemy.orm 

14from slixmpp.exceptions import XMPPError 

15from slixmpp.plugins.xep_0231.stanza import BitsOfBinary 

16from sqlalchemy import ColumnElement, Engine, delete, event, select, update 

17from sqlalchemy.exc import InvalidRequestError 

18from sqlalchemy.orm import ( 

19 Session, 

20 attributes, 

21 joinedload, 

22 load_only, 

23 selectinload, 

24 sessionmaker, 

25 with_loader_criteria, 

26) 

27 

28from ..core import config 

29from ..util.archive_msg import HistoryMessage 

30from ..util.types import MamMetadata, Sticker 

31from .models import ( 

32 ArchivedMessage, 

33 ArchivedMessageSource, 

34 Attachment, 

35 Avatar, 

36 Bob, 

37 Contact, 

38 ContactSent, 

39 DirectMessages, 

40 DirectThreads, 

41 GatewayUser, 

42 GroupMessages, 

43 GroupMessagesOrigin, 

44 GroupThreads, 

45 Participant, 

46 Room, 

47 Space, 

48) 

49 

50 

51class UpdatableBase(Protocol): 

52 id: ClassVar[ColumnElement[int]] 

53 user_account_id: ClassVar[ColumnElement[int]] 

54 updated: ClassVar[ColumnElement[bool]] 

55 

56 

57T = TypeVar("T", bound=UpdatableBase) 

58 

59 

60class UpdatedMixin[T]: 

61 model: type[T] = NotImplemented 

62 

63 def __init__(self, session: Session) -> None: 

64 self.reset_updated(session) 

65 

66 def get_by_pk(self, session: Session, pk: int) -> T | None: 

67 stmt = select(self.model).where(self.model.id == pk) # type:ignore[attr-defined] 

68 return session.scalar(stmt) 

69 

70 def reset_updated(self, session: Session) -> None: 

71 session.execute(update(self.model).values(updated=False)) 

72 

73 def get_for(self, session: Session, user_pk: int) -> list[T]: 

74 stmt = select(self.model).where(self.model.user_account_id == user_pk) # type:ignore[attr-defined] 

75 return list(session.scalars(stmt)) 

76 

77 

78class SlidgeStore: 

79 def __init__(self, engine: Engine) -> None: 

80 self._engine = engine 

81 self.session = sessionmaker[Any](engine) 

82 

83 self.users = UserStore(self.session) 

84 self.avatars = AvatarStore(self.session) 

85 self.id_map = IdMapStore() 

86 self.bob = BobStore() 

87 self.attachments = AttachmentStore() 

88 with self.session() as session: 

89 self.contacts = ContactStore(session) 

90 self.mam = MAMStore(session, self.session) 

91 self.rooms = RoomStore(session) 

92 self.participants = ParticipantStore(session) 

93 self.spaces = SpaceStore(session) 

94 session.commit() 

95 

96 

97class UserStore: 

98 def __init__(self, session_maker: sessionmaker[Any]) -> None: 

99 self.session = session_maker 

100 

101 def update(self, user: GatewayUser) -> None: 

102 with self.session(expire_on_commit=False) as session: 

103 # https://github.com/sqlalchemy/sqlalchemy/discussions/6473 

104 try: 

105 attributes.flag_modified(user, "legacy_module_data") 

106 attributes.flag_modified(user, "preferences") 

107 except InvalidRequestError: 

108 pass 

109 session.add(user) 

110 session.commit() 

111 

112 

113class AvatarStore: 

114 def __init__(self, session_maker: sessionmaker[Any]) -> None: 

115 self.session = session_maker 

116 

117 

118LegacyToXmppType = ( 

119 type[DirectMessages] 

120 | type[DirectThreads] 

121 | type[GroupMessages] 

122 | type[GroupThreads] 

123 | type[GroupMessagesOrigin] 

124) 

125 

126 

127class IdMapStore: 

128 @staticmethod 

129 def _set( 

130 session: Session, 

131 foreign_key: int, 

132 legacy_id: str, 

133 xmpp_ids: list[str], 

134 type_: LegacyToXmppType, 

135 ) -> None: 

136 kwargs = {"foreign_key": foreign_key, "legacy_id": legacy_id} 

137 ids = list( 

138 session.scalars( 

139 select(type_.id).filter( 

140 type_.foreign_key == foreign_key, type_.legacy_id == legacy_id 

141 ) 

142 ) 

143 ) 

144 if ids: 

145 log.debug("Resetting legacy ID %s", legacy_id) 

146 session.execute(delete(type_).where(type_.id.in_(ids))) 

147 for xmpp_id in xmpp_ids: 

148 msg = type_(xmpp_id=xmpp_id, **kwargs) 

149 session.add(msg) 

150 

151 def set_thread( 

152 self, 

153 session: Session, 

154 foreign_key: int, 

155 legacy_id: str, 

156 xmpp_id: str, 

157 group: bool, 

158 ) -> None: 

159 self._set( 

160 session, 

161 foreign_key, 

162 legacy_id, 

163 [xmpp_id], 

164 GroupThreads if group else DirectThreads, 

165 ) 

166 

167 def set_msg( 

168 self, 

169 session: Session, 

170 foreign_key: int, 

171 legacy_id: str, 

172 xmpp_ids: list[str], 

173 group: bool, 

174 ) -> None: 

175 self._set( 

176 session, 

177 foreign_key, 

178 legacy_id, 

179 xmpp_ids, 

180 GroupMessages if group else DirectMessages, 

181 ) 

182 

183 def set_origin( 

184 self, session: Session, foreign_key: int, legacy_id: str, xmpp_id: str 

185 ) -> None: 

186 self._set( 

187 session, 

188 foreign_key, 

189 legacy_id, 

190 [xmpp_id], 

191 GroupMessagesOrigin, 

192 ) 

193 

194 def get_origin( 

195 self, session: Session, foreign_key: int, legacy_id: str 

196 ) -> list[str]: 

197 return self._get( 

198 session, 

199 foreign_key, 

200 legacy_id, 

201 GroupMessagesOrigin, 

202 ) 

203 

204 @staticmethod 

205 def _get( 

206 session: Session, foreign_key: int, legacy_id: str, type_: LegacyToXmppType 

207 ) -> list[str]: 

208 return list( 

209 session.scalars( 

210 select(type_.xmpp_id).filter_by( 

211 foreign_key=foreign_key, legacy_id=str(legacy_id) 

212 ) 

213 ) 

214 ) 

215 

216 def get_xmpp( 

217 self, session: Session, foreign_key: int, legacy_id: str, group: bool 

218 ) -> list[str]: 

219 return self._get( 

220 session, 

221 foreign_key, 

222 legacy_id, 

223 GroupMessages if group else DirectMessages, 

224 ) 

225 

226 @staticmethod 

227 def _get_legacy( 

228 session: Session, foreign_key: int, xmpp_id: str, type_: LegacyToXmppType 

229 ) -> str | None: 

230 return session.scalar( 

231 select(type_.legacy_id).filter_by(foreign_key=foreign_key, xmpp_id=xmpp_id) 

232 ) 

233 

234 def get_legacy( 

235 self, 

236 session: Session, 

237 foreign_key: int, 

238 xmpp_id: str, 

239 group: bool, 

240 origin: bool = False, 

241 ) -> str | None: 

242 if origin and group: 

243 return self._get_legacy( 

244 session, 

245 foreign_key, 

246 xmpp_id, 

247 GroupMessagesOrigin, 

248 ) 

249 return self._get_legacy( 

250 session, 

251 foreign_key, 

252 xmpp_id, 

253 GroupMessages if group else DirectMessages, 

254 ) 

255 

256 def get_thread( 

257 self, session: Session, foreign_key: int, xmpp_id: str, group: bool 

258 ) -> str | None: 

259 return self._get_legacy( 

260 session, 

261 foreign_key, 

262 xmpp_id, 

263 GroupThreads if group else DirectThreads, 

264 ) 

265 

266 @staticmethod 

267 def was_sent_by_user( 

268 session: Session, foreign_key: int, legacy_id: str, group: bool 

269 ) -> bool: 

270 type_ = GroupMessages if group else DirectMessages 

271 return ( 

272 session.scalar( 

273 select(type_.id).filter_by(foreign_key=foreign_key, legacy_id=legacy_id) 

274 ) 

275 is not None 

276 ) 

277 

278 

279class ContactStore(UpdatedMixin[Contact]): 

280 model = Contact 

281 

282 def __init__(self, session: Session) -> None: 

283 super().__init__(session) 

284 session.execute(update(Contact).values(cached_presence=False)) 

285 session.execute(update(Contact).values(caps_ver=None)) 

286 

287 @staticmethod 

288 def add_to_sent(session: Session, contact_pk: int, msg_id: str) -> None: 

289 if ( 

290 session.query(ContactSent.id) 

291 .where(ContactSent.contact_id == contact_pk) 

292 .where(ContactSent.msg_id == msg_id) 

293 .first() 

294 ) is not None: 

295 log.warning("Contact %s has already sent message %s", contact_pk, msg_id) 

296 return 

297 new = ContactSent(contact_id=contact_pk, msg_id=msg_id) 

298 session.add(new) 

299 

300 @staticmethod 

301 def pop_sent_up_to(session: Session, contact_pk: int, msg_id: str) -> list[str]: 

302 result = [] 

303 to_del = [] 

304 for row in session.execute( 

305 select(ContactSent) 

306 .where(ContactSent.contact_id == contact_pk) 

307 .order_by(ContactSent.id) 

308 ).scalars(): 

309 to_del.append(row.id) 

310 result.append(row.msg_id) 

311 if row.msg_id == msg_id: 

312 break 

313 session.execute(delete(ContactSent).where(ContactSent.id.in_(to_del))) 

314 return result 

315 

316 

317class MAMStore: 

318 def __init__(self, session: Session, session_maker: sessionmaker[Any]) -> None: 

319 self.session = session_maker 

320 self.reset_source(session) 

321 

322 @staticmethod 

323 def reset_source(session: Session) -> None: 

324 session.execute( 

325 update(ArchivedMessage).values(source=ArchivedMessageSource.BACKFILL) 

326 ) 

327 

328 @staticmethod 

329 def nuke_older_than(session: Session, days: int) -> None: 

330 session.execute( 

331 delete(ArchivedMessage).where( 

332 ArchivedMessage.timestamp < datetime.now(tz=UTC) - timedelta(days=days) 

333 ) 

334 ) 

335 

336 @staticmethod 

337 def add_message( 

338 session: Session, 

339 room_pk: int, 

340 message: HistoryMessage, 

341 archive_only: bool, 

342 legacy_msg_id: str | None, 

343 ) -> None: 

344 source = ( 

345 ArchivedMessageSource.BACKFILL 

346 if archive_only 

347 else ArchivedMessageSource.LIVE 

348 ) 

349 existing = session.execute( 

350 select(ArchivedMessage) 

351 .where(ArchivedMessage.room_id == room_pk) 

352 .where(ArchivedMessage.stanza_id == message.id) 

353 ).scalar() 

354 if existing is None and legacy_msg_id is not None: 

355 existing = session.execute( 

356 select(ArchivedMessage) 

357 .where(ArchivedMessage.room_id == room_pk) 

358 .where(ArchivedMessage.legacy_id == str(legacy_msg_id)) 

359 ).scalar() 

360 if existing is not None: 

361 log.debug("Updating message %s in room %s", message.id, room_pk) 

362 existing.timestamp = message.when 

363 existing.stanza = str(message.stanza) 

364 existing.author_jid = message.stanza.get_from() 

365 existing.source = source 

366 existing.legacy_id = legacy_msg_id 

367 session.add(existing) 

368 return 

369 mam_msg = ArchivedMessage( 

370 stanza_id=message.id, 

371 timestamp=message.when, 

372 stanza=str(message.stanza), 

373 author_jid=message.stanza.get_from(), 

374 room_id=room_pk, 

375 source=source, 

376 legacy_id=legacy_msg_id, 

377 ) 

378 session.add(mam_msg) 

379 

380 @staticmethod 

381 def get_messages( 

382 session: Session, 

383 room_pk: int, 

384 start_date: datetime | None = None, 

385 end_date: datetime | None = None, 

386 before_id: str | None = None, 

387 after_id: str | None = None, 

388 ids: Collection[str] = (), 

389 last_page_n: int | None = None, 

390 sender: str | None = None, 

391 flip: bool = False, 

392 ) -> Iterator[HistoryMessage]: 

393 q = select(ArchivedMessage).where(ArchivedMessage.room_id == room_pk) 

394 if start_date is not None: 

395 q = q.where(ArchivedMessage.timestamp >= start_date) 

396 if end_date is not None: 

397 q = q.where(ArchivedMessage.timestamp <= end_date) 

398 if before_id is not None: 

399 stamp = session.execute( 

400 select(ArchivedMessage.timestamp).where( 

401 ArchivedMessage.stanza_id == before_id, 

402 ArchivedMessage.room_id == room_pk, 

403 ) 

404 ).scalar_one_or_none() 

405 if stamp is None: 

406 raise XMPPError( 

407 "item-not-found", 

408 f"Message {before_id} not found", 

409 ) 

410 q = q.where(ArchivedMessage.timestamp < stamp) 

411 if after_id is not None: 

412 stamp = session.execute( 

413 select(ArchivedMessage.timestamp).where( 

414 ArchivedMessage.stanza_id == after_id, 

415 ArchivedMessage.room_id == room_pk, 

416 ) 

417 ).scalar_one_or_none() 

418 if stamp is None: 

419 raise XMPPError( 

420 "item-not-found", 

421 f"Message {after_id} not found", 

422 ) 

423 q = q.where(ArchivedMessage.timestamp > stamp) 

424 if ids: 

425 q = q.filter(ArchivedMessage.stanza_id.in_(ids)) 

426 if sender is not None: 

427 q = q.where(ArchivedMessage.author_jid == sender) 

428 if flip: 

429 q = q.order_by(ArchivedMessage.timestamp.desc()) 

430 else: 

431 q = q.order_by(ArchivedMessage.timestamp.asc()) 

432 msgs = list(session.execute(q).scalars()) 

433 if ids and len(msgs) != len(ids): 

434 raise XMPPError( 

435 "item-not-found", 

436 "One of the requested messages IDs could not be found " 

437 "with the given constraints.", 

438 ) 

439 if last_page_n is not None: 

440 msgs = msgs[:last_page_n] if flip else msgs[-last_page_n:] 

441 for h in msgs: 

442 yield HistoryMessage( 

443 stanza=str(h.stanza), when=h.timestamp.replace(tzinfo=UTC) 

444 ) 

445 

446 @staticmethod 

447 def get_first( 

448 session: Session, room_pk: int, with_legacy_id: bool = False 

449 ) -> ArchivedMessage | None: 

450 q = ( 

451 select(ArchivedMessage) 

452 .where(ArchivedMessage.room_id == room_pk) 

453 .order_by(ArchivedMessage.timestamp.asc()) 

454 ) 

455 if with_legacy_id: 

456 q = q.filter(ArchivedMessage.legacy_id.isnot(None)) 

457 return session.execute(q).scalar() 

458 

459 @staticmethod 

460 def get_last( 

461 session: Session, room_pk: int, source: ArchivedMessageSource | None = None 

462 ) -> ArchivedMessage | None: 

463 q = select(ArchivedMessage).where(ArchivedMessage.room_id == room_pk) 

464 

465 if source is not None: 

466 q = q.where(ArchivedMessage.source == source) 

467 

468 return session.execute(q.order_by(ArchivedMessage.timestamp.desc())).scalar() 

469 

470 def get_first_and_last(self, session: Session, room_pk: int) -> list[MamMetadata]: 

471 r = [] 

472 first = self.get_first(session, room_pk) 

473 if first is not None: 

474 r.append(MamMetadata(first.stanza_id, first.timestamp)) 

475 last = self.get_last(session, room_pk) 

476 if last is not None: 

477 r.append(MamMetadata(last.stanza_id, last.timestamp)) 

478 return r 

479 

480 @staticmethod 

481 def get_most_recent_with_legacy_id( 

482 session: Session, room_pk: int, source: ArchivedMessageSource | None = None 

483 ) -> ArchivedMessage | None: 

484 q = ( 

485 select(ArchivedMessage) 

486 .where(ArchivedMessage.room_id == room_pk) 

487 .where(ArchivedMessage.legacy_id.isnot(None)) 

488 ) 

489 if source is not None: 

490 q = q.where(ArchivedMessage.source == source) 

491 return session.execute(q.order_by(ArchivedMessage.timestamp.desc())).scalar() 

492 

493 @staticmethod 

494 def get_least_recent_with_legacy_id_after( 

495 session: Session, 

496 room_pk: int, 

497 after_id: str, 

498 source: ArchivedMessageSource = ArchivedMessageSource.LIVE, 

499 ) -> ArchivedMessage | None: 

500 after_timestamp = ( 

501 session.query(ArchivedMessage.timestamp) 

502 .filter(ArchivedMessage.room_id == room_pk) 

503 .filter(ArchivedMessage.legacy_id == after_id) 

504 .scalar() 

505 ) 

506 q = ( 

507 select(ArchivedMessage) 

508 .where(ArchivedMessage.room_id == room_pk) 

509 .where(ArchivedMessage.legacy_id.isnot(None)) 

510 .where(ArchivedMessage.source == source) 

511 .where(ArchivedMessage.timestamp > after_timestamp) 

512 ) 

513 return session.execute(q.order_by(ArchivedMessage.timestamp.asc())).scalar() 

514 

515 @staticmethod 

516 def get_by_legacy_id( 

517 session: Session, room_pk: int, legacy_id: str 

518 ) -> ArchivedMessage | None: 

519 return ( 

520 session.query(ArchivedMessage) 

521 .filter(ArchivedMessage.room_id == room_pk) 

522 .filter(ArchivedMessage.legacy_id == legacy_id) 

523 .first() 

524 ) 

525 

526 @staticmethod 

527 def pop_unread_up_to(session: Session, room_pk: int, stanza_id: str) -> list[str]: 

528 q = ( 

529 select(ArchivedMessage.id, ArchivedMessage.stanza_id) 

530 .where(ArchivedMessage.room_id == room_pk) 

531 .where(~ArchivedMessage.displayed_by_user) 

532 .where(ArchivedMessage.legacy_id.is_not(None)) 

533 .order_by(ArchivedMessage.timestamp.asc()) 

534 ) 

535 

536 ref = session.scalar( 

537 select(ArchivedMessage) 

538 .where(ArchivedMessage.room_id == room_pk) 

539 .where(ArchivedMessage.stanza_id == stanza_id) 

540 ) 

541 

542 if ref is None: 

543 log.debug( 

544 "(pop unread in muc): message not found, returning all MAM messages." 

545 ) 

546 rows = session.execute(q) 

547 else: 

548 rows = session.execute(q.where(ArchivedMessage.timestamp <= ref.timestamp)) 

549 

550 pks: list[int] = [] 

551 stanza_ids: list[str] = [] 

552 

553 for id_, sid in rows: 

554 pks.append(id_) 

555 stanza_ids.append(sid) 

556 

557 session.execute( 

558 update(ArchivedMessage) 

559 .where(ArchivedMessage.id.in_(pks)) 

560 .values(displayed_by_user=True) 

561 ) 

562 return stanza_ids 

563 

564 @staticmethod 

565 def is_displayed_by_user( 

566 session: Session, room_jid_localpart: str, legacy_msg_id: str 

567 ) -> bool: 

568 return any( 

569 session.execute( 

570 select(ArchivedMessage.displayed_by_user) 

571 .join(Room) 

572 .where(Room.jid_localpart == room_jid_localpart) 

573 .where(ArchivedMessage.legacy_id == legacy_msg_id) 

574 ).scalars() 

575 ) 

576 

577 

578class RoomStore(UpdatedMixin[Room]): 

579 model = Room 

580 

581 def reset_updated(self, session: Session) -> None: 

582 super().reset_updated(session) 

583 session.execute( 

584 update(Room).values( 

585 subject_setter=None, 

586 user_resources=None, 

587 history_filled=False, 

588 participants_filled=False, 

589 ) 

590 ) 

591 

592 @staticmethod 

593 def get_all(session: Session, user_pk: int) -> Iterator[Room]: 

594 yield from session.scalars(select(Room).where(Room.user_account_id == user_pk)) 

595 

596 @staticmethod 

597 def get(session: Session, user_pk: int, legacy_id: str) -> Room: 

598 return session.execute( 

599 select(Room) 

600 .where(Room.user_account_id == user_pk) 

601 .where(Room.legacy_id == legacy_id) 

602 ).scalar_one() 

603 

604 @staticmethod 

605 def nick_available(session: Session, room_pk: int, nickname: str) -> bool: 

606 return ( 

607 session.execute( 

608 select(Participant.id).filter_by(room_id=room_pk, nickname=nickname) 

609 ) 

610 ).one_or_none() is None 

611 

612 

613class ParticipantStore: 

614 def __init__(self, session: Session) -> None: 

615 session.execute(delete(Participant)) 

616 

617 @staticmethod 

618 def get_all( 

619 session: Session, room_pk: int, user_included: bool = True 

620 ) -> Iterator[Participant]: 

621 query = select(Participant).where(Participant.room_id == room_pk) 

622 if not user_included: 

623 query = query.where(~Participant.is_user) 

624 yield from session.scalars(query).unique() 

625 

626 @staticmethod 

627 def delete(session: Session, pk: int) -> None: 

628 session.execute(delete(Participant).where(Participant.id == pk)) 

629 

630 

631class BobStore: 

632 _ATTR_MAP: ClassVar[dict[str, str]] = { 

633 "sha-1": "sha_1", 

634 "sha1": "sha_1", 

635 "sha-256": "sha_256", 

636 "sha256": "sha_256", 

637 "sha-512": "sha_512", 

638 "sha512": "sha_512", 

639 } 

640 

641 _ALG_MAP: ClassVar[dict[str, Callable[[bytes], hashlib._Hash]]] = { 

642 "sha_1": hashlib.sha1, 

643 "sha_256": hashlib.sha256, 

644 "sha_512": hashlib.sha512, 

645 } 

646 

647 def __init__(self) -> None: 

648 if (config.HOME_DIR / "slidge_stickers").exists(): 

649 shutil.move( 

650 config.HOME_DIR / "slidge_stickers", config.HOME_DIR / "bob_store" 

651 ) 

652 self.root_dir = config.HOME_DIR / "bob_store" 

653 self.root_dir.mkdir(exist_ok=True) 

654 

655 @staticmethod 

656 def __split_cid(cid: str) -> list[str]: 

657 return cid.removesuffix("@bob.xmpp.org").split("+") 

658 

659 def __get_condition(self, cid: str) -> ColumnElement[bool]: 

660 alg_name, digest = self.__split_cid(cid) 

661 attr = self._ATTR_MAP.get(alg_name) 

662 if attr is None: 

663 log.warning("Unknown hash algorithm: %s", alg_name) 

664 raise ValueError 

665 return getattr(Bob, attr) == digest # type:ignore[no-any-return] 

666 

667 def get(self, session: Session, cid: str) -> Bob | None: 

668 try: 

669 return session.query(Bob).filter(self.__get_condition(cid)).scalar() # type:ignore[no-any-return] 

670 except ValueError: 

671 log.warning("Cannot get Bob with CID: %s", cid) 

672 return None 

673 

674 def get_sticker(self, session: Session, cid: str) -> Sticker | None: 

675 bob = self.get(session, cid) 

676 if bob is None: 

677 return None 

678 return self.__sticker_from_bob(bob) 

679 

680 def __sticker_from_bob(self, bob: Bob) -> Sticker: 

681 return Sticker( 

682 self.root_dir / bob.file_name, 

683 bob.content_type, 

684 {h: getattr(bob, h) for h in self._ALG_MAP}, 

685 ) 

686 

687 def get_bob( 

688 self, session: Session, _jid: object, _node: object, _ifrom: object, cid: str 

689 ) -> BitsOfBinary | None: 

690 stored = self.get(session, cid) 

691 if stored is None: 

692 return None 

693 bob = BitsOfBinary() 

694 bob["data"] = (self.root_dir / stored.file_name).read_bytes() 

695 if stored.content_type is not None: 

696 bob["type"] = stored.content_type 

697 bob["cid"] = cid 

698 return bob 

699 

700 def del_bob( 

701 self, session: Session, _jid: object, _node: object, _ifrom: object, cid: str 

702 ) -> None: 

703 try: 

704 file_name = session.scalar( 

705 delete(Bob).where(self.__get_condition(cid)).returning(Bob.file_name) 

706 ) 

707 except ValueError: 

708 log.warning("Cannot delete Bob with CID: %s", cid) 

709 return 

710 if file_name is None: 

711 log.warning("No BoB with CID: %s", cid) 

712 return 

713 (self.root_dir / file_name).unlink() 

714 

715 def set_bob( 

716 self, 

717 session: Session, 

718 _jid: object, 

719 _node: object, 

720 _ifrom: object, 

721 bob: BitsOfBinary, 

722 ) -> Sticker | None: 

723 return self.set_sticker(session, bob["cid"], bob["data"], bob["type"]) 

724 

725 def set_sticker( 

726 self, 

727 session: Session, 

728 cid: str, 

729 bytes_: bytes, 

730 content_type: str | None, 

731 ) -> Sticker | None: 

732 try: 

733 alg_name, digest = self.__split_cid(cid) 

734 except ValueError: 

735 log.warning("Invalid CID provided: %s", cid) 

736 return None 

737 attr = self._ATTR_MAP.get(alg_name) 

738 if attr is None: 

739 log.warning("Cannot set Bob: Unknown algorithm type: %s", alg_name) 

740 return None 

741 existing = self.get(session, cid) 

742 if existing: 

743 log.debug("Bob already exists") 

744 return None 

745 path = self.root_dir / uuid.uuid4().hex 

746 if content_type is None: 

747 try: 

748 import magic 

749 except ImportError: 

750 content_type = "application/octet-stream" 

751 else: 

752 content_type = magic.from_buffer(bytes_, mime=True) 

753 path = path.with_suffix(guess_extension(content_type) or "") 

754 path.write_bytes(bytes_) 

755 hashes = {k: v(bytes_).hexdigest() for k, v in self._ALG_MAP.items()} 

756 if hashes[attr] != digest: 

757 path.unlink(missing_ok=True) 

758 raise ValueError("Provided CID does not match calculated hash") 

759 row = Bob(file_name=path.name, content_type=content_type, **hashes) 

760 session.add(row) 

761 return self.__sticker_from_bob(row) 

762 

763 

764class SpaceStore(UpdatedMixin[Space]): 

765 model = Space 

766 

767 def __init__(self, session: Session) -> None: 

768 session.execute(delete(Space)) 

769 

770 @staticmethod 

771 def add_or_get(session: Session, user_pk: int, legacy_id: str) -> Space: 

772 space = session.execute( 

773 select(Space) 

774 .where(Space.user_account_id == user_pk) 

775 .where(Space.legacy_id == legacy_id) 

776 .options( 

777 joinedload(Space.avatar), 

778 joinedload(Space.banner), 

779 ) 

780 ).scalar_one_or_none() 

781 if space is None: 

782 space = Space( 

783 user_account_id=user_pk, 

784 legacy_id=legacy_id, 

785 avatar=None, 

786 banner=None, 

787 rooms=[], 

788 ) 

789 session.add(space) 

790 session.commit() 

791 return space 

792 

793 @staticmethod 

794 def get_all(session: Session, user_pk: int) -> Iterable[Space]: 

795 return session.execute( 

796 select(Space).where(Space.user_account_id == user_pk) 

797 ).scalars() 

798 

799 @staticmethod 

800 def get_by_legacy_id( 

801 session: Session, 

802 user_pk: int, 

803 legacy_id: str, 

804 affiliations: bool = False, 

805 images: bool = False, 

806 room_legacy_id_filter: Iterable[str] | None = None, 

807 ) -> Space | None: 

808 stmt = ( 

809 select(Space) 

810 .where(Space.user_account_id == user_pk) 

811 .where(Space.legacy_id == legacy_id) 

812 ) 

813 if affiliations: 

814 stmt = stmt.options( 

815 joinedload(Space.owners), 

816 joinedload(Space.creator), 

817 ) 

818 if images: 

819 stmt = stmt.options( 

820 joinedload(Space.avatar), 

821 joinedload(Space.banner), 

822 ) 

823 if room_legacy_id_filter is not None: 

824 stmt = stmt.options( 

825 selectinload(Space.rooms), 

826 with_loader_criteria( 

827 Room, 

828 Room.legacy_id.in_(room_legacy_id_filter), 

829 ), 

830 ) 

831 

832 return session.execute(stmt).unique().scalar_one_or_none() 

833 

834 @staticmethod 

835 def get_unupdated(session: Session, user_pk: int) -> list[Space]: 

836 return list( 

837 session.execute( 

838 select(Space) 

839 .where(Space.user_account_id == user_pk) 

840 .where(Space.updated.is_(False)) 

841 .options( 

842 joinedload(Space.avatar), 

843 joinedload(Space.banner), 

844 ) 

845 ).scalars() 

846 ) 

847 

848 @staticmethod 

849 def get_rooms( 

850 session: Session, 

851 user_pk: int, 

852 legacy_id: str, 

853 room_legacy_ids: Iterable[str] = (), 

854 ) -> list[Room]: 

855 q = ( 

856 select(Room) 

857 .join(Room.space) 

858 .where(Room.user_account_id == user_pk) 

859 .where(Space.legacy_id == legacy_id) 

860 .options(load_only(Room.jid_localpart, Room.name)) 

861 ) 

862 if room_legacy_ids: 

863 q = q.where(Room.legacy_id.in_(room_legacy_ids)) 

864 return list(session.execute(q).scalars()) 

865 

866 @staticmethod 

867 def exists(session: Session, user_pk: int, legacy_id: str) -> bool: 

868 return session.execute( 

869 select( 

870 sa.exists() 

871 .where(Space.user_account_id == user_pk) 

872 .where(Space.legacy_id == legacy_id) 

873 ) 

874 ).scalar_one() 

875 

876 

877class AttachmentStore: 

878 @staticmethod 

879 def get_all(session: sa.orm.Session) -> list[Attachment]: 

880 return list( 

881 session.execute( 

882 select(Attachment).options(load_only(Attachment.id, Attachment.url)) 

883 ).scalars() 

884 ) 

885 

886 @staticmethod 

887 def remove(session: sa.orm.Session, pks: list[int]) -> None: 

888 session.execute(delete(Attachment).where(Attachment.id.in_(pks))) 

889 

890 

891@event.listens_for(sa.orm.Session, "after_flush") 

892def _check_avatar_orphans(session: Session, flush_context: sa.ExecutionContext) -> None: 

893 if not session.deleted: 

894 return 

895 

896 potentially_orphaned = set() 

897 for obj in session.deleted: 

898 if isinstance(obj, (Contact, Room)) and obj.avatar_id: 

899 potentially_orphaned.add(obj.avatar_id) 

900 if not potentially_orphaned: 

901 return 

902 

903 result = session.execute( 

904 sa.delete(Avatar).where( 

905 sa.and_( 

906 Avatar.id.in_(potentially_orphaned), 

907 sa.not_(sa.exists().where(Contact.avatar_id == Avatar.id)), 

908 sa.not_(sa.exists().where(Room.avatar_id == Avatar.id)), 

909 sa.not_(sa.exists().where(Space.avatar_id == Avatar.id)), 

910 sa.not_(sa.exists().where(Space.banner_id == Avatar.id)), 

911 ) 

912 ) 

913 ) 

914 deleted_count = result.rowcount # type:ignore[attr-defined] 

915 log.debug("Auto-deleted %s orphaned avatars", deleted_count) 

916 

917 

918log = logging.getLogger(__name__)