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
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-10 04:45 +0000
1from __future__ import annotations
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
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)
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)
51class UpdatableBase(Protocol):
52 id: ClassVar[ColumnElement[int]]
53 user_account_id: ClassVar[ColumnElement[int]]
54 updated: ClassVar[ColumnElement[bool]]
57T = TypeVar("T", bound=UpdatableBase)
60class UpdatedMixin[T]:
61 model: type[T] = NotImplemented
63 def __init__(self, session: Session) -> None:
64 self.reset_updated(session)
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)
70 def reset_updated(self, session: Session) -> None:
71 session.execute(update(self.model).values(updated=False))
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))
78class SlidgeStore:
79 def __init__(self, engine: Engine) -> None:
80 self._engine = engine
81 self.session = sessionmaker[Any](engine)
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()
97class UserStore:
98 def __init__(self, session_maker: sessionmaker[Any]) -> None:
99 self.session = session_maker
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()
113class AvatarStore:
114 def __init__(self, session_maker: sessionmaker[Any]) -> None:
115 self.session = session_maker
118LegacyToXmppType = (
119 type[DirectMessages]
120 | type[DirectThreads]
121 | type[GroupMessages]
122 | type[GroupThreads]
123 | type[GroupMessagesOrigin]
124)
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)
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 )
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 )
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 )
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 )
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 )
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 )
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 )
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 )
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 )
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 )
279class ContactStore(UpdatedMixin[Contact]):
280 model = Contact
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))
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)
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
317class MAMStore:
318 def __init__(self, session: Session, session_maker: sessionmaker[Any]) -> None:
319 self.session = session_maker
320 self.reset_source(session)
322 @staticmethod
323 def reset_source(session: Session) -> None:
324 session.execute(
325 update(ArchivedMessage).values(source=ArchivedMessageSource.BACKFILL)
326 )
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 )
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)
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 )
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()
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)
465 if source is not None:
466 q = q.where(ArchivedMessage.source == source)
468 return session.execute(q.order_by(ArchivedMessage.timestamp.desc())).scalar()
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
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()
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()
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 )
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 )
536 ref = session.scalar(
537 select(ArchivedMessage)
538 .where(ArchivedMessage.room_id == room_pk)
539 .where(ArchivedMessage.stanza_id == stanza_id)
540 )
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))
550 pks: list[int] = []
551 stanza_ids: list[str] = []
553 for id_, sid in rows:
554 pks.append(id_)
555 stanza_ids.append(sid)
557 session.execute(
558 update(ArchivedMessage)
559 .where(ArchivedMessage.id.in_(pks))
560 .values(displayed_by_user=True)
561 )
562 return stanza_ids
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 )
578class RoomStore(UpdatedMixin[Room]):
579 model = Room
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 )
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))
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()
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
613class ParticipantStore:
614 def __init__(self, session: Session) -> None:
615 session.execute(delete(Participant))
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()
626 @staticmethod
627 def delete(session: Session, pk: int) -> None:
628 session.execute(delete(Participant).where(Participant.id == pk))
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 }
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 }
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)
655 @staticmethod
656 def __split_cid(cid: str) -> list[str]:
657 return cid.removesuffix("@bob.xmpp.org").split("+")
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]
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
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)
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 )
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
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()
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"])
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)
764class SpaceStore(UpdatedMixin[Space]):
765 model = Space
767 def __init__(self, session: Session) -> None:
768 session.execute(delete(Space))
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
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()
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 )
832 return session.execute(stmt).unique().scalar_one_or_none()
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 )
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())
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()
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 )
886 @staticmethod
887 def remove(session: sa.orm.Session, pks: list[int]) -> None:
888 session.execute(delete(Attachment).where(Attachment.id.in_(pks)))
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
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
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)
918log = logging.getLogger(__name__)