Coverage for slidge/db/store.py: 86%
477 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-10 12:26 +0000
1from __future__ import annotations
3import hashlib
4import json
5import logging
6import shutil
7import uuid
8from collections.abc import Callable, Collection, Iterable, Iterator
9from datetime import UTC, datetime, timedelta
10from mimetypes import guess_extension
11from typing import Any, ClassVar, Protocol, TypeVar
13import sqlalchemy as sa
14import sqlalchemy.orm
15from slixmpp import Message
16from slixmpp.exceptions import XMPPError
17from slixmpp.plugins.xep_0231.stanza import BitsOfBinary
18from slixmpp.xmlstream import tostring
19from sqlalchemy import ColumnElement, Engine, delete, event, select, update
20from sqlalchemy.exc import InvalidRequestError
21from sqlalchemy.orm import (
22 Session,
23 attributes,
24 joinedload,
25 load_only,
26 selectinload,
27 sessionmaker,
28 with_loader_criteria,
29)
31from ..core import config
32from ..util.types import MamMetadata, Sticker
33from ..util.util import uniquify_replace
34from .models import (
35 ArchivedMessage,
36 ArchivedMessageSource,
37 Attachment,
38 Avatar,
39 Bob,
40 Contact,
41 ContactSent,
42 DirectMessages,
43 DirectThreads,
44 GatewayUser,
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]]
55 legacy_id: ClassVar[ColumnElement[str]]
58T = TypeVar("T", bound=UpdatableBase)
61class UpdatedMixin[T: Contact | Room | Participant | Space]:
62 model: type[T] = NotImplemented
64 def get_by_pk(self, session: Session, pk: int) -> T | None:
65 stmt = select(self.model).where(self.model.id == pk)
66 return session.scalar(stmt)
68 def get_for(self, session: Session, user_pk: int) -> list[T]:
69 stmt = select(self.model).where(self.model.user_account_id == user_pk) # type:ignore[attr-defined]
70 return list(session.scalars(stmt))
72 def by_legacy_id(self, session: Session, user_pk: int, legacy_id: str) -> T | None:
73 return session.scalar(
74 select(self.model).where( # type:ignore[arg-type]
75 self.model.user_account_id == user_pk, # type:ignore[attr-defined]
76 self.model.legacy_id == legacy_id, # type:ignore[attr-defined]
77 )
78 )
81class SlidgeStore:
82 def __init__(self, engine: Engine) -> None:
83 self._engine = engine
84 self.session = sessionmaker[Any](engine)
86 self.users = UserStore(self.session)
87 self.avatars = AvatarStore(self.session)
88 self.id_map = IdMapStore()
89 self.bob = BobStore()
90 self.attachments = AttachmentStore()
91 self.participants = ParticipantStore()
92 self.spaces = SpaceStore()
93 with self.session() as session:
94 self.contacts = ContactStore(session)
95 self.mam = MAMStore(session, self.session)
96 self.rooms = RoomStore(session)
97 session.commit()
100class UserStore:
101 def __init__(self, session_maker: sessionmaker[Any]) -> None:
102 self.session = session_maker
104 def update(self, user: GatewayUser) -> None:
105 with self.session(expire_on_commit=False) as session:
106 # https://github.com/sqlalchemy/sqlalchemy/discussions/6473
107 try:
108 attributes.flag_modified(user, "legacy_module_data")
109 attributes.flag_modified(user, "preferences")
110 except InvalidRequestError:
111 pass
112 session.add(user)
113 session.commit()
116class AvatarStore:
117 def __init__(self, session_maker: sessionmaker[Any]) -> None:
118 self.session = session_maker
121LegacyToXmppType = type[DirectMessages] | type[DirectThreads]
124class IdMapStore:
125 @staticmethod
126 def _set(
127 session: Session,
128 foreign_key: int,
129 legacy_id: str,
130 xmpp_ids: Iterable[str],
131 type_: LegacyToXmppType,
132 ) -> None:
133 kwargs = {"foreign_key": foreign_key, "legacy_id": legacy_id}
134 ids = list(
135 session.scalars(
136 select(type_.id).filter(
137 type_.foreign_key == foreign_key, type_.legacy_id == legacy_id
138 )
139 )
140 )
141 if ids:
142 log.debug("Resetting legacy ID %s", legacy_id)
143 session.execute(delete(type_).where(type_.id.in_(ids)))
144 for xmpp_id in xmpp_ids:
145 msg = type_(xmpp_id=xmpp_id, **kwargs)
146 session.add(msg)
148 def set_thread(
149 self,
150 session: Session,
151 foreign_key: int,
152 legacy_id: str,
153 xmpp_id: str,
154 group: bool,
155 ) -> None:
156 if group:
157 if xmpp_id == legacy_id:
158 # No need to store in this case. The thread column of the mam
159 # table is already populated, and if thread_legacy_id is NULL,
160 # we assume xmpp_id == legacy_id.
161 return
162 session.execute(
163 update(ArchivedMessage)
164 .where(
165 ArchivedMessage.room_id == foreign_key,
166 ArchivedMessage.stanza_id == xmpp_id,
167 )
168 .values(thread_legacy_id=legacy_id)
169 )
170 else:
171 self._set(session, foreign_key, legacy_id, [xmpp_id], DirectThreads)
173 def set_msg(
174 self,
175 session: Session,
176 foreign_key: int,
177 legacy_id: str,
178 xmpp_ids: Iterable[str],
179 group: bool,
180 ) -> None:
181 if group:
182 session.execute(
183 update(ArchivedMessage)
184 .where(
185 ArchivedMessage.room_id == foreign_key,
186 ArchivedMessage.stanza_id.in_(xmpp_ids),
187 )
188 .values(legacy_id=legacy_id)
189 )
190 else:
191 # In direct messages, we need to keep the mapping metween XMPP
192 # client-generated msg-id and legacy ID, and we also need to know
193 # which messages have been sent by the user.
194 self._set(session, foreign_key, legacy_id, xmpp_ids, DirectMessages)
196 def set_origin(
197 self, session: Session, foreign_key: int, legacy_id: str, xmpp_id: str
198 ) -> None:
199 session.execute(
200 update(ArchivedMessage)
201 .where(
202 ArchivedMessage.room_id == foreign_key,
203 ArchivedMessage.legacy_id == legacy_id,
204 )
205 .values(origin_id=xmpp_id)
206 )
208 def get_origin(
209 self, session: Session, foreign_key: int, legacy_id: str
210 ) -> list[str]:
211 # it's unlikely type checkers will ever be able to determine that this
212 # is filtered out for NULL values
213 return list( # ty:ignore[invalid-return-type]
214 session.execute( # type:ignore[arg-type]
215 select(ArchivedMessage.origin_id).where(
216 ArchivedMessage.room_id == foreign_key,
217 ArchivedMessage.legacy_id == legacy_id,
218 ArchivedMessage.origin_id.is_not(None),
219 )
220 ).scalars()
221 )
223 @staticmethod
224 def _get(
225 session: Session, foreign_key: int, legacy_id: str, type_: LegacyToXmppType
226 ) -> list[str]:
227 return list(
228 session.scalars(
229 select(type_.xmpp_id).filter_by(
230 foreign_key=foreign_key, legacy_id=str(legacy_id)
231 )
232 )
233 )
235 def get_xmpp(
236 self, session: Session, foreign_key: int, legacy_id: str, group: bool
237 ) -> list[str]:
238 if group:
239 return list(
240 session.execute(
241 select(ArchivedMessage.stanza_id).where(
242 ArchivedMessage.room_id == foreign_key,
243 ArchivedMessage.legacy_id == legacy_id,
244 )
245 ).scalars()
246 )
247 else:
248 return self._get(session, foreign_key, legacy_id, DirectMessages)
250 @staticmethod
251 def _get_legacy(
252 session: Session, foreign_key: int, xmpp_id: str, type_: LegacyToXmppType
253 ) -> str | None:
254 return session.scalar(
255 select(type_.legacy_id).filter_by(foreign_key=foreign_key, xmpp_id=xmpp_id)
256 )
258 def get_legacy(
259 self,
260 session: Session,
261 foreign_key: int,
262 xmpp_id: str,
263 group: bool,
264 origin: bool = False,
265 ) -> str | None:
266 if group:
267 cond = (
268 (ArchivedMessage.origin_id == xmpp_id)
269 if origin
270 else (ArchivedMessage.stanza_id == xmpp_id)
271 )
272 return session.execute(
273 select(ArchivedMessage.legacy_id).where(
274 ArchivedMessage.room_id == foreign_key,
275 cond,
276 )
277 ).scalar_one_or_none()
279 return self._get_legacy(session, foreign_key, xmpp_id, DirectMessages)
281 def get_thread(
282 self, session: Session, foreign_key: int, xmpp_id: str, group: bool
283 ) -> str | None:
284 if group:
285 row = session.execute(
286 select(
287 ArchivedMessage.thread,
288 ArchivedMessage.thread_legacy_id,
289 ).where(
290 ArchivedMessage.thread == xmpp_id,
291 )
292 ).first()
293 if row is None:
294 return None
295 xmpp_id_, legacy_id = row
296 return legacy_id or xmpp_id_
297 else:
298 return self._get_legacy(
299 session,
300 foreign_key,
301 xmpp_id,
302 DirectThreads,
303 )
305 @staticmethod
306 def was_sent_by_user(session: Session, foreign_key: int, legacy_id: str) -> bool:
307 """
308 Only works for direct messages
309 """
310 return (
311 session.scalar(
312 select(DirectMessages.id).filter_by(
313 foreign_key=foreign_key, legacy_id=legacy_id
314 )
315 )
316 is not None
317 )
320class ContactStore(UpdatedMixin[Contact]):
321 model = Contact
323 def __init__(self, session: Session) -> None:
324 super().__init__()
326 @staticmethod
327 def add_to_sent(session: Session, contact_pk: int, msg_id: str) -> None:
328 if (
329 session.query(ContactSent.id)
330 .where(ContactSent.contact_id == contact_pk)
331 .where(ContactSent.msg_id == msg_id)
332 .first()
333 ) is not None:
334 log.warning("Contact %s has already sent message %s", contact_pk, msg_id)
335 return
336 new = ContactSent(contact_id=contact_pk, msg_id=msg_id)
337 session.add(new)
339 @staticmethod
340 def pop_sent_up_to(session: Session, contact_pk: int, msg_id: str) -> list[str]:
341 result = []
342 to_del = []
343 for row in session.execute(
344 select(ContactSent)
345 .where(ContactSent.contact_id == contact_pk)
346 .order_by(ContactSent.id)
347 ).scalars():
348 to_del.append(row.id)
349 result.append(row.msg_id)
350 if row.msg_id == msg_id:
351 break
352 session.execute(delete(ContactSent).where(ContactSent.id.in_(to_del)))
353 return result
356class MAMStore:
357 def __init__(self, session: Session, session_maker: sessionmaker[Any]) -> None:
358 self.session = session_maker
359 self.reset_source(session)
361 @staticmethod
362 def reset_source(session: Session) -> None:
363 session.execute(
364 update(ArchivedMessage).values(source=ArchivedMessageSource.BACKFILL)
365 )
367 @staticmethod
368 def nuke_older_than(session: Session, days: int) -> None:
369 session.execute(
370 delete(ArchivedMessage).where(
371 ArchivedMessage.timestamp < datetime.now(tz=UTC) - timedelta(days=days)
372 )
373 )
375 @staticmethod
376 def add_message(
377 session: Session,
378 room_pk: int,
379 stanza: Message,
380 backfill: bool,
381 legacy_msg_id: str | None,
382 from_user: bool = False,
383 ) -> ArchivedMessage:
384 source = (
385 ArchivedMessageSource.BACKFILL if backfill else ArchivedMessageSource.LIVE
386 )
388 when = stanza["delay"]["stamp"] or datetime.now(tz=UTC)
389 stanza_id = stanza["stanza_id"]["id"]
390 assert stanza_id
392 origin_id = stanza["origin_id"]["id"] or None
393 author_resource = stanza.get_from().resource
394 author_affiliation = stanza["muc"]["affiliation"] or None
395 author_role = stanza["muc"]["role"] or None
397 if from_user:
398 author_jid_localpart = None
399 else:
400 author_jid = stanza["muc"]["jid"]
401 author_jid_localpart = author_jid.user if author_jid else None
403 occupant_id = stanza["occupant-id"]["id"]
405 displayed_id = stanza["displayed"]["id"] if "displayed" in stanza else None
407 if "reactions" in stanza:
408 reaction_id = stanza["reactions"]["id"]
409 if reaction_id is None:
410 log.warning("Received a reaction without an ID?: %r", stanza)
411 reaction_id = reaction_emojis = None
412 else:
413 # list() because this returns a set, which would require some more
414 # SQLAlchemy plumbing. A list is fine!
415 reaction_emojis = list(stanza["reactions"]["values"]) or None
416 else:
417 reaction_id = reaction_emojis = None
419 body = stanza["body"]
420 thread = stanza["thread"]
422 del stanza["delay"]
423 del stanza["markable"]
424 del stanza["store"]
425 del stanza["chat_state"]
426 del stanza["muc"]
427 del stanza["origin_id"]
428 del stanza["stanza_id"]
429 del stanza["occupant-id"]
430 del stanza["displayed"]
431 del stanza["reactions"]
432 del stanza["body"]
433 del stanza["thread"]
435 payload = _serialize_children(stanza) or None
437 existing = session.execute(
438 select(ArchivedMessage).where(
439 ArchivedMessage.room_id == room_pk,
440 ArchivedMessage.stanza_id == stanza_id,
441 )
442 ).scalar()
443 if existing is None and legacy_msg_id is not None:
444 existing = session.execute(
445 select(ArchivedMessage).where(
446 ArchivedMessage.room_id == room_pk,
447 ArchivedMessage.legacy_id == legacy_msg_id,
448 )
449 ).scalar()
450 elif payload is None and displayed_id is not None and occupant_id is not None:
451 existing_displayed_marker = session.scalar(
452 select(ArchivedMessage).where(
453 ArchivedMessage.room_id == room_pk,
454 ArchivedMessage.occupant_id == occupant_id,
455 ArchivedMessage.displayed_id == displayed_id,
456 )
457 )
458 if existing_displayed_marker is not None:
459 log.debug("Received a duplicate displayed marker, ignoring it")
460 return existing_displayed_marker
462 if existing is not None:
463 log.warning(
464 "Updating message %s in room %s. This should not happen",
465 stanza_id,
466 room_pk,
467 )
468 stanza_id = uniquify_replace(stanza_id)
470 mam_msg = ArchivedMessage(stanza_id=stanza_id, room_id=room_pk)
472 mam_msg.stanza_id = stanza_id
473 mam_msg.timestamp = when
474 mam_msg.payload = payload
475 mam_msg.origin_id = origin_id
476 mam_msg.author_resource = author_resource
477 mam_msg.author_affiliation = author_affiliation
478 mam_msg.author_role = author_role
479 mam_msg.author_jid_localpart = author_jid_localpart
480 mam_msg.room_id = room_pk
481 mam_msg.source = source
482 mam_msg.legacy_id = legacy_msg_id
483 mam_msg.occupant_id = occupant_id
484 mam_msg.from_user = from_user
485 mam_msg.displayed_id = displayed_id
486 mam_msg.reaction_emojis = reaction_emojis
487 mam_msg.reaction_id = reaction_id
488 mam_msg.body = body
489 mam_msg.thread = thread
491 session.add(mam_msg)
493 return mam_msg
495 @staticmethod
496 def get_messages(
497 session: Session,
498 room_pk: int,
499 start_date: datetime | None = None,
500 end_date: datetime | None = None,
501 before_id: str | None = None,
502 after_id: str | None = None,
503 ids: Collection[str] = (),
504 last_page_n: int | None = None,
505 author_resource: str | None = None,
506 flip: bool = False,
507 ) -> Iterator[ArchivedMessage]:
508 q = select(ArchivedMessage).where(ArchivedMessage.room_id == room_pk)
509 if start_date is not None:
510 q = q.where(ArchivedMessage.timestamp >= start_date)
511 if end_date is not None:
512 q = q.where(ArchivedMessage.timestamp <= end_date)
513 if before_id is not None:
514 stamp = session.execute(
515 select(ArchivedMessage.timestamp).where(
516 ArchivedMessage.stanza_id == before_id,
517 ArchivedMessage.room_id == room_pk,
518 )
519 ).scalar_one_or_none()
520 if stamp is None:
521 raise XMPPError(
522 "item-not-found",
523 f"Message {before_id} not found",
524 )
525 q = q.where(ArchivedMessage.timestamp < stamp)
526 if after_id is not None:
527 stamp = session.execute(
528 select(ArchivedMessage.timestamp).where(
529 ArchivedMessage.stanza_id == after_id,
530 ArchivedMessage.room_id == room_pk,
531 )
532 ).scalar_one_or_none()
533 if stamp is None:
534 raise XMPPError(
535 "item-not-found",
536 f"Message {after_id} not found",
537 )
538 q = q.where(ArchivedMessage.timestamp > stamp)
539 if ids:
540 q = q.filter(ArchivedMessage.stanza_id.in_(ids))
541 if author_resource is not None:
542 q = q.where(ArchivedMessage.author_resource == author_resource)
543 if flip:
544 q = q.order_by(ArchivedMessage.timestamp.desc())
545 else:
546 q = q.order_by(ArchivedMessage.timestamp.asc())
547 msgs = list(session.execute(q).scalars())
548 if ids and len(msgs) != len(ids):
549 raise XMPPError(
550 "item-not-found",
551 "One of the requested messages IDs could not be found "
552 "with the given constraints.",
553 )
554 if last_page_n is not None:
555 msgs = msgs[:last_page_n] if flip else msgs[-last_page_n:]
556 yield from msgs
558 @staticmethod
559 def get_first(
560 session: Session, room_pk: int, with_legacy_id: bool = False
561 ) -> ArchivedMessage | None:
562 q = (
563 select(ArchivedMessage)
564 .where(ArchivedMessage.room_id == room_pk)
565 .order_by(ArchivedMessage.timestamp.asc())
566 )
567 if with_legacy_id:
568 q = q.where(ArchivedMessage.legacy_id.is_not(None))
569 return session.scalar(q.limit(1))
571 @staticmethod
572 def get_last(
573 session: Session, room_pk: int, source: ArchivedMessageSource | None = None
574 ) -> ArchivedMessage | None:
575 q = select(ArchivedMessage).where(ArchivedMessage.room_id == room_pk)
576 if source is not None:
577 q = q.where(ArchivedMessage.source == source)
578 return session.scalar(q.order_by(ArchivedMessage.timestamp.desc()).limit(1))
580 def get_first_and_last(self, session: Session, room_pk: int) -> list[MamMetadata]:
581 r = []
582 first = self.get_first(session, room_pk)
583 if first is not None:
584 r.append(MamMetadata(first.stanza_id, first.timestamp))
585 last = self.get_last(session, room_pk)
586 if last is not None:
587 r.append(MamMetadata(last.stanza_id, last.timestamp))
588 return r
590 @staticmethod
591 def get_most_recent_with_legacy_id(
592 session: Session, room_pk: int, source: ArchivedMessageSource | None = None
593 ) -> ArchivedMessage | None:
594 q = (
595 select(ArchivedMessage)
596 .where(ArchivedMessage.room_id == room_pk)
597 .where(ArchivedMessage.legacy_id.isnot(None))
598 )
599 if source is not None:
600 q = q.where(ArchivedMessage.source == source)
601 return session.scalar(q.order_by(ArchivedMessage.timestamp.desc()).limit(1))
603 @staticmethod
604 def get_least_recent_with_legacy_id_after(
605 session: Session,
606 room_pk: int,
607 after_id: str,
608 source: ArchivedMessageSource = ArchivedMessageSource.LIVE,
609 ) -> ArchivedMessage | None:
610 return session.scalar(
611 select(ArchivedMessage)
612 .where(
613 ArchivedMessage.room_id == room_pk,
614 ArchivedMessage.legacy_id.isnot(None),
615 ArchivedMessage.source == source,
616 ArchivedMessage.timestamp
617 > (
618 select(sa.func.max(ArchivedMessage.timestamp))
619 .where(
620 ArchivedMessage.room_id == room_pk,
621 ArchivedMessage.legacy_id == after_id,
622 )
623 .scalar_subquery()
624 ),
625 )
626 .order_by(ArchivedMessage.timestamp.asc(), ArchivedMessage.id.desc())
627 .limit(1)
628 )
630 @staticmethod
631 def get_by_legacy_id(
632 session: Session, room_pk: int, legacy_id: str
633 ) -> ArchivedMessage | None:
634 return (
635 session.query(ArchivedMessage)
636 .filter(ArchivedMessage.room_id == room_pk)
637 .filter(ArchivedMessage.legacy_id == legacy_id)
638 .first()
639 )
641 @staticmethod
642 def pop_unread_up_to(session: Session, room_pk: int, stanza_id: str) -> list[str]:
643 q = (
644 select(ArchivedMessage.id, ArchivedMessage.stanza_id)
645 .where(ArchivedMessage.room_id == room_pk)
646 .where(~ArchivedMessage.displayed_by_user)
647 .where(ArchivedMessage.legacy_id.is_not(None))
648 .order_by(ArchivedMessage.timestamp.asc())
649 )
651 ref = session.scalar(
652 select(ArchivedMessage)
653 .where(ArchivedMessage.room_id == room_pk)
654 .where(ArchivedMessage.stanza_id == stanza_id)
655 )
657 if ref is None:
658 log.debug(
659 "(pop unread in muc): message not found, returning all MAM messages."
660 )
661 rows = session.execute(q)
662 else:
663 rows = session.execute(q.where(ArchivedMessage.timestamp <= ref.timestamp))
665 pks: list[int] = []
666 stanza_ids: list[str] = []
668 for id_, sid in rows:
669 pks.append(id_)
670 stanza_ids.append(sid)
672 session.execute(
673 update(ArchivedMessage)
674 .where(ArchivedMessage.id.in_(pks))
675 .values(displayed_by_user=True)
676 )
677 return stanza_ids
679 @staticmethod
680 def is_displayed_by_user(
681 session: Session, room_jid_localpart: str, legacy_msg_id: str
682 ) -> bool:
683 return any(
684 session.execute(
685 select(ArchivedMessage.displayed_by_user)
686 .join(Room)
687 .where(Room.jid_localpart == room_jid_localpart)
688 .where(ArchivedMessage.legacy_id == legacy_msg_id)
689 ).scalars()
690 )
692 @staticmethod
693 def get_body(session: Session, room_pk: int, legacy_msg_id: str) -> str | None:
694 return session.scalar(
695 select(ArchivedMessage.body).where(
696 ArchivedMessage.room_id == room_pk,
697 ArchivedMessage.legacy_id == legacy_msg_id,
698 )
699 )
702class RoomStore(UpdatedMixin[Room]):
703 model = Room
705 def __init__(self, session: Session) -> None:
706 session.execute(
707 update(Room).values(
708 user_resources=None,
709 history_filled=False,
710 )
711 )
713 @staticmethod
714 def get_all(session: Session, user_pk: int) -> Iterator[Room]:
715 yield from session.scalars(select(Room).where(Room.user_account_id == user_pk))
717 @staticmethod
718 def get(session: Session, user_pk: int, legacy_id: str) -> Room:
719 return session.execute(
720 select(Room)
721 .where(Room.user_account_id == user_pk)
722 .where(Room.legacy_id == legacy_id)
723 ).scalar_one()
725 @staticmethod
726 def resource_available(
727 session: Session, room_pk: int, resource: str, exclude_pk: int | None = None
728 ) -> bool:
729 q = select(Participant.id).where(
730 Participant.room_id == room_pk,
731 (Participant.nickname == resource) | (Participant._resource == resource),
732 )
733 if exclude_pk is not None:
734 q = q.where(Participant.id != exclude_pk)
736 return (session.scalar(q)) is None
739class ParticipantStore:
740 @staticmethod
741 def get_all(
742 session: Session, room_pk: int, user_included: bool = True
743 ) -> Iterator[Participant]:
744 query = select(Participant).where(Participant.room_id == room_pk)
745 if not user_included:
746 query = query.where(~Participant.is_user)
747 yield from session.scalars(query).unique()
749 @staticmethod
750 def get_user_participants(
751 session: Session, user_pk: int
752 ) -> Iterable[tuple[str, list[str], str]]:
753 for res, nick, resources, localpart in session.execute(
754 sa.select(
755 Participant._resource,
756 Participant.nickname,
757 Room.user_resources,
758 Room.jid_localpart,
759 )
760 .join(Room, Participant.room_id == Room.id)
761 .where(Room.user_account_id == user_pk)
762 .where(Room.user_resources.isnot(None))
763 .where(Participant.is_user)
764 ):
765 assert resources is not None
766 yield res or nick, json.loads(resources), localpart
768 @staticmethod
769 def delete(session: Session, pk: int) -> None:
770 session.execute(delete(Participant).where(Participant.id == pk))
773class BobStore:
774 _ATTR_MAP: ClassVar[dict[str, str]] = {
775 "sha-1": "sha_1",
776 "sha1": "sha_1",
777 "sha-256": "sha_256",
778 "sha256": "sha_256",
779 "sha-512": "sha_512",
780 "sha512": "sha_512",
781 }
783 _ALG_MAP: ClassVar[dict[str, Callable[[bytes], hashlib._Hash]]] = {
784 "sha_1": hashlib.sha1,
785 "sha_256": hashlib.sha256,
786 "sha_512": hashlib.sha512,
787 }
789 def __init__(self) -> None:
790 if (config.HOME_DIR / "slidge_stickers").exists():
791 shutil.move(
792 config.HOME_DIR / "slidge_stickers", config.HOME_DIR / "bob_store"
793 )
794 self.root_dir = config.HOME_DIR / "bob_store"
795 self.root_dir.mkdir(exist_ok=True)
797 @staticmethod
798 def __split_cid(cid: str) -> list[str]:
799 return cid.removesuffix("@bob.xmpp.org").split("+")
801 def __get_condition(self, cid: str) -> ColumnElement[bool]:
802 alg_name, digest = self.__split_cid(cid)
803 attr = self._ATTR_MAP.get(alg_name)
804 if attr is None:
805 log.warning("Unknown hash algorithm: %s", alg_name)
806 raise ValueError
807 return getattr(Bob, attr) == digest # type:ignore[no-any-return]
809 def get(self, session: Session, cid: str) -> Bob | None:
810 try:
811 return session.query(Bob).filter(self.__get_condition(cid)).scalar() # type:ignore[no-any-return]
812 except ValueError:
813 log.warning("Cannot get Bob with CID: %s", cid)
814 return None
816 def get_sticker(self, session: Session, cid: str) -> Sticker | None:
817 bob = self.get(session, cid)
818 if bob is None:
819 return None
820 return self.__sticker_from_bob(bob)
822 def __sticker_from_bob(self, bob: Bob) -> Sticker:
823 return Sticker(
824 self.root_dir / bob.file_name,
825 bob.content_type,
826 {h: getattr(bob, h) for h in self._ALG_MAP},
827 )
829 def get_bob(
830 self, session: Session, _jid: object, _node: object, _ifrom: object, cid: str
831 ) -> BitsOfBinary | None:
832 stored = self.get(session, cid)
833 if stored is None:
834 return None
835 bob = BitsOfBinary()
836 bob["data"] = (self.root_dir / stored.file_name).read_bytes()
837 if stored.content_type is not None:
838 bob["type"] = stored.content_type
839 bob["cid"] = cid
840 return bob
842 def del_bob(
843 self, session: Session, _jid: object, _node: object, _ifrom: object, cid: str
844 ) -> None:
845 try:
846 file_name = session.scalar(
847 delete(Bob).where(self.__get_condition(cid)).returning(Bob.file_name)
848 )
849 except ValueError:
850 log.warning("Cannot delete Bob with CID: %s", cid)
851 return
852 if file_name is None:
853 log.warning("No BoB with CID: %s", cid)
854 return
855 (self.root_dir / file_name).unlink()
857 def set_bob(
858 self,
859 session: Session,
860 _jid: object,
861 _node: object,
862 _ifrom: object,
863 bob: BitsOfBinary,
864 ) -> Sticker | None:
865 return self.set_sticker(session, bob["cid"], bob["data"], bob["type"])
867 def set_sticker(
868 self,
869 session: Session,
870 cid: str,
871 bytes_: bytes,
872 content_type: str | None,
873 ) -> Sticker | None:
874 try:
875 alg_name, digest = self.__split_cid(cid)
876 except ValueError:
877 log.warning("Invalid CID provided: %s", cid)
878 return None
879 attr = self._ATTR_MAP.get(alg_name)
880 if attr is None:
881 log.warning("Cannot set Bob: Unknown algorithm type: %s", alg_name)
882 return None
883 existing = self.get(session, cid)
884 if existing:
885 log.debug("Bob already exists")
886 return None
887 path = self.root_dir / uuid.uuid4().hex
888 if content_type is None:
889 try:
890 import magic
891 except ImportError:
892 content_type = "application/octet-stream"
893 else:
894 content_type = magic.from_buffer(bytes_, mime=True)
895 path = path.with_suffix(guess_extension(content_type) or "")
896 path.write_bytes(bytes_)
897 hashes = {k: v(bytes_).hexdigest() for k, v in self._ALG_MAP.items()}
898 if hashes[attr] != digest:
899 path.unlink(missing_ok=True)
900 raise ValueError("Provided CID does not match calculated hash")
901 row = Bob(file_name=path.name, content_type=content_type, **hashes)
902 session.add(row)
903 return self.__sticker_from_bob(row)
906class SpaceStore(UpdatedMixin[Space]):
907 model = Space
909 @staticmethod
910 def add_or_get(session: Session, user_pk: int, legacy_id: str) -> Space:
911 space = session.execute(
912 select(Space)
913 .where(Space.user_account_id == user_pk)
914 .where(Space.legacy_id == legacy_id)
915 .options(
916 joinedload(Space.avatar),
917 joinedload(Space.banner),
918 )
919 ).scalar_one_or_none()
920 if space is None:
921 space = Space(
922 user_account_id=user_pk,
923 legacy_id=legacy_id,
924 avatar=None,
925 banner=None,
926 rooms=[],
927 )
928 session.add(space)
929 session.commit()
930 return space
932 @staticmethod
933 def get_all(session: Session, user_pk: int) -> Iterable[Space]:
934 return session.execute(
935 select(Space).where(Space.user_account_id == user_pk)
936 ).scalars()
938 @staticmethod
939 def get_by_legacy_id(
940 session: Session,
941 user_pk: int,
942 legacy_id: str,
943 affiliations: bool = False,
944 images: bool = False,
945 room_legacy_id_filter: Iterable[str] | None = None,
946 ) -> Space | None:
947 stmt = (
948 select(Space)
949 .where(Space.user_account_id == user_pk)
950 .where(Space.legacy_id == legacy_id)
951 )
952 if affiliations:
953 stmt = stmt.options(
954 joinedload(Space.owners),
955 joinedload(Space.creator),
956 )
957 if images:
958 stmt = stmt.options(
959 joinedload(Space.avatar),
960 joinedload(Space.banner),
961 )
962 if room_legacy_id_filter is not None:
963 stmt = stmt.options(
964 selectinload(Space.rooms),
965 with_loader_criteria(
966 Room,
967 Room.legacy_id.in_(room_legacy_id_filter),
968 ),
969 )
971 return session.execute(stmt).unique().scalar_one_or_none()
973 @staticmethod
974 def get_rooms(
975 session: Session,
976 user_pk: int,
977 legacy_id: str,
978 room_legacy_ids: Iterable[str] = (),
979 ) -> list[Room]:
980 q = (
981 select(Room)
982 .join(Room.space)
983 .where(Room.user_account_id == user_pk)
984 .where(Space.legacy_id == legacy_id)
985 .options(load_only(Room.jid_localpart, Room.name))
986 )
987 if room_legacy_ids:
988 q = q.where(Room.legacy_id.in_(room_legacy_ids))
989 return list(session.execute(q).scalars())
991 @staticmethod
992 def exists(session: Session, user_pk: int, legacy_id: str) -> bool:
993 return session.execute(
994 select(
995 sa.exists()
996 .where(Space.user_account_id == user_pk)
997 .where(Space.legacy_id == legacy_id)
998 )
999 ).scalar_one()
1002class AttachmentStore:
1003 @staticmethod
1004 def get_all(session: sa.orm.Session) -> list[Attachment]:
1005 return list(
1006 session.execute(
1007 select(Attachment).options(load_only(Attachment.id, Attachment.url))
1008 ).scalars()
1009 )
1011 @staticmethod
1012 def remove(session: sa.orm.Session, pks: list[int]) -> None:
1013 session.execute(delete(Attachment).where(Attachment.id.in_(pks)))
1016@event.listens_for(sa.orm.Session, "after_flush")
1017def _check_avatar_orphans(session: Session, flush_context: sa.ExecutionContext) -> None:
1018 if not session.deleted:
1019 return
1021 potentially_orphaned = set()
1022 for obj in session.deleted:
1023 if isinstance(obj, (Contact, Room)) and obj.avatar_id:
1024 potentially_orphaned.add(obj.avatar_id)
1025 if not potentially_orphaned:
1026 return
1028 result = session.execute(
1029 sa.delete(Avatar).where(
1030 sa.and_(
1031 Avatar.id.in_(potentially_orphaned),
1032 sa.not_(sa.exists().where(Contact.avatar_id == Avatar.id)),
1033 sa.not_(sa.exists().where(Room.avatar_id == Avatar.id)),
1034 sa.not_(sa.exists().where(Space.avatar_id == Avatar.id)),
1035 sa.not_(sa.exists().where(Space.banner_id == Avatar.id)),
1036 )
1037 )
1038 )
1039 deleted_count = result.rowcount # type:ignore[attr-defined]
1040 log.debug("Auto-deleted %s orphaned avatars", deleted_count)
1043def _serialize_children(msg: Message) -> str:
1044 return "".join(tostring(child, xmlns=msg.namespace) for child in msg.xml)
1047log = logging.getLogger(__name__)