slidge 0.2.0a0__py3-none-any.whl → 0.2.0a1__py3-none-any.whl

slidge/db/store.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import json
 import logging
-import warnings
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Collection, Iterator, Optional, Type
@@ -11,6 +10,7 @@ from slixmpp import JID, Iq, Message, Presence
 from slixmpp.exceptions import XMPPError
 from sqlalchemy import Engine, delete, select, update
 from sqlalchemy.orm import Session, attributes
+from sqlalchemy.sql.functions import count
 
 from ..util.archive_msg import HistoryMessage
 from ..util.types import URL, CachedPresence
@@ -19,6 +19,7 @@ from ..util.types import MamMetadata, MucAffiliation, MucRole
 from .meta import Base
 from .models import (
     ArchivedMessage,
+    ArchivedMessageSource,
     Attachment,
     Avatar,
     Contact,
@@ -63,9 +64,18 @@ class UpdatedMixin(EngineMixin):
     def __init__(self, *a, **kw):
         super().__init__(*a, **kw)
         with self.session() as session:
+            session.execute(
+                delete(self.model).where(~self.model.updated)  # type:ignore
+            )
             session.execute(update(self.model).values(updated=False))  # type:ignore
             session.commit()
 
+    def get_by_pk(self, pk: int) -> Optional[Base]:
+        with self.session() as session:
+            return session.execute(
+                select(self.model).where(self.model.id == pk)  # type:ignore
+            ).scalar()
+
 
 class SlidgeStore(EngineMixin):
     def __init__(self, engine: Engine):
@@ -134,34 +144,18 @@ class AvatarStore(EngineMixin):
                 select(Avatar).where(Avatar.legacy_id == legacy_id)
             ).scalar()
 
-    def store_jid(self, sha: str, jid: JID) -> None:
+    def get_by_pk(self, pk: int) -> Optional[Avatar]:
         with self.session() as session:
-            pk = session.execute(select(Avatar.id).where(Avatar.hash == sha)).scalar()
-            if pk is None:
-                warnings.warn("avatar not found")
-                return
-            session.execute(
-                update(Contact).where(Contact.jid == jid.bare).values(avatar_id=pk)
-            )
-            session.execute(
-                update(Room).where(Room.jid == jid.bare).values(avatar_id=pk)
-            )
-            session.commit()
+            return session.execute(select(Avatar).where(Avatar.id == pk)).scalar()
 
-    def get_by_jid(self, jid: JID) -> Optional[Avatar]:
+    def delete_by_pk(self, pk: int):
         with self.session() as session:
-            avatar = session.execute(
-                select(Avatar).where(Avatar.contacts.any(Contact.jid == jid))
-            ).scalar()
-            if avatar is not None:
-                return avatar
-            return session.execute(
-                select(Avatar).where(Avatar.rooms.any(Room.jid == jid))
-            ).scalar()
+            session.execute(delete(Avatar).where(Avatar.id == pk))
+            session.commit()
 
-    def get_by_pk(self, pk: int) -> Optional[Avatar]:
+    def get_all(self) -> Iterator[Avatar]:
         with self.session() as session:
-            return session.execute(select(Avatar).where(Avatar.id == pk)).scalar()
+            yield from session.execute(select(Avatar)).scalars()
 
 
 class SentStore(EngineMixin):
@@ -273,22 +267,6 @@ class ContactStore(UpdatedMixin):
             session.execute(update(Contact).values(cached_presence=False))
             session.commit()
 
-    def add(self, user_pk: int, legacy_id: str, contact_jid: JID) -> int:
-        with self.session() as session:
-            existing = session.execute(
-                select(Contact)
-                .where(Contact.legacy_id == legacy_id)
-                .where(Contact.jid == contact_jid.bare)
-            ).scalar()
-            if existing is not None:
-                return existing.id
-            contact = Contact(
-                jid=contact_jid.bare, legacy_id=legacy_id, user_account_id=user_pk
-            )
-            session.add(contact)
-            session.commit()
-            return contact.id
-
     def get_all(self, user_pk: int) -> Iterator[Contact]:
         with self.session() as session:
             yield from session.execute(
@@ -375,20 +353,36 @@ class ContactStore(UpdatedMixin):
                 return None
             return contact.avatar.legacy_id
 
-    def update(self, contact: "LegacyContact"):
+    def update(self, contact: "LegacyContact", commit=True) -> int:
         with self.session() as session:
-            session.execute(
-                update(Contact)
-                .where(Contact.id == contact.contact_pk)
-                .values(
-                    nick=contact.name,
-                    is_friend=contact.is_friend,
-                    added_to_roster=contact.added_to_roster,
-                    updated=True,
-                    extra_attributes=contact.serialize_extra_attributes(),
+            if contact.contact_pk is None:
+                if contact.cached_presence is not None:
+                    presence_kwargs = contact.cached_presence._asdict()
+                    presence_kwargs["cached_presence"] = True
+                else:
+                    presence_kwargs = {}
+                row = Contact(
+                    jid=contact.jid.bare,
+                    legacy_id=str(contact.legacy_id),
+                    user_account_id=contact.user_pk,
+                    **presence_kwargs,
                 )
-            )
-            session.commit()
+            else:
+                row = (
+                    session.query(Contact)
+                    .filter(Contact.id == contact.contact_pk)
+                    .one()
+                )
+            row.nick = contact.name
+            row.is_friend = contact.is_friend
+            row.added_to_roster = contact.added_to_roster
+            row.updated = True
+            row.extra_attributes = contact.serialize_extra_attributes()
+            row.caps_ver = contact._caps_ver
+            session.add(row)
+            if commit:
+                session.commit()
+            return row.id
 
     def add_to_sent(self, contact_pk: int, msg_id: str) -> None:
         with self.session() as session:
@@ -422,19 +416,52 @@ class ContactStore(UpdatedMixin):
             )
             session.commit()
 
+    def set_added_to_roster(self, contact_pk: int, value: bool) -> None:
+        with self.session() as session:
+            session.execute(
+                update(Contact)
+                .where(Contact.id == contact_pk)
+                .values(added_to_roster=value)
+            )
+            session.commit()
+
+    def delete(self, contact_pk: int) -> None:
+        with self.session() as session:
+            session.execute(delete(Contact).where(Contact.id == contact_pk))
+            session.commit()
+
 
 class MAMStore(EngineMixin):
+    def __init__(self, *a, **kw):
+        super().__init__(*a, **kw)
+        with self.session() as session:
+            session.execute(
+                update(ArchivedMessage).values(source=ArchivedMessageSource.BACKFILL)
+            )
+            session.commit()
+
     def nuke_older_than(self, days: int) -> None:
         with self.session() as session:
             session.execute(
                 delete(ArchivedMessage).where(
-                    ArchivedMessage.timestamp > datetime.now() - timedelta(days=days)
+                    ArchivedMessage.timestamp < datetime.now() - timedelta(days=days)
                 )
             )
             session.commit()
 
-    def add_message(self, room_pk: int, message: HistoryMessage) -> None:
+    def add_message(
+        self,
+        room_pk: int,
+        message: HistoryMessage,
+        archive_only: bool,
+        legacy_msg_id: str | None,
+    ) -> None:
         with self.session() as session:
+            source = (
+                ArchivedMessageSource.BACKFILL
+                if archive_only
+                else ArchivedMessageSource.LIVE
+            )
             existing = session.execute(
                 select(ArchivedMessage)
                 .where(ArchivedMessage.room_id == room_pk)
@@ -445,6 +472,8 @@ class MAMStore(EngineMixin):
                 existing.timestamp = message.when
                 existing.stanza = str(message.stanza)
                 existing.author_jid = message.stanza.get_from()
+                existing.source = source
+                existing.legacy_id = legacy_msg_id
                 session.add(existing)
                 session.commit()
                 return
@@ -454,6 +483,8 @@ class MAMStore(EngineMixin):
                 stanza=str(message.stanza),
                 author_jid=message.stanza.get_from(),
                 room_id=room_pk,
+                source=source,
+                legacy_id=legacy_msg_id,
             )
             session.add(mam_msg)
             session.commit()
@@ -526,25 +557,83 @@ class MAMStore(EngineMixin):
                     stanza=str(h.stanza), when=h.timestamp.replace(tzinfo=timezone.utc)
                 )
 
-    def get_first_and_last(self, room_pk: int) -> list[MamMetadata]:
-        r = []
+    def get_first(self, room_pk: int, with_legacy_id=False) -> ArchivedMessage | None:
         with self.session() as session:
-            first = session.execute(
-                select(ArchivedMessage.stanza_id, ArchivedMessage.timestamp)
+            q = (
+                select(ArchivedMessage)
                 .where(ArchivedMessage.room_id == room_pk)
                 .order_by(ArchivedMessage.timestamp.asc())
-            ).first()
+            )
+            if with_legacy_id:
+                q = q.filter(ArchivedMessage.legacy_id.isnot(None))
+            return session.execute(q).scalar()
+
+    def get_last(
+        self, room_pk: int, source: ArchivedMessageSource | None = None
+    ) -> ArchivedMessage | None:
+        with self.session() as session:
+            q = select(ArchivedMessage).where(ArchivedMessage.room_id == room_pk)
+
+            if source is not None:
+                q = q.where(ArchivedMessage.source == source)
+
+            return session.execute(
+                q.order_by(ArchivedMessage.timestamp.desc())
+            ).scalar()
+
+    def get_first_and_last(self, room_pk: int) -> list[MamMetadata]:
+        r = []
+        with self.session():
+            first = self.get_first(room_pk)
             if first is not None:
-                r.append(MamMetadata(*first))
-            last = session.execute(
-                select(ArchivedMessage.stanza_id, ArchivedMessage.timestamp)
-                .where(ArchivedMessage.room_id == room_pk)
-                .order_by(ArchivedMessage.timestamp.desc())
-            ).first()
+                r.append(MamMetadata(first.stanza_id, first.timestamp))
+            last = self.get_last(room_pk)
             if last is not None:
-                r.append(MamMetadata(*last))
+                r.append(MamMetadata(last.stanza_id, last.timestamp))
         return r
 
+    def get_most_recent_with_legacy_id(
+        self, room_pk: int, source: ArchivedMessageSource | None = None
+    ) -> ArchivedMessage | None:
+        with self.session() as session:
+            q = (
+                select(ArchivedMessage)
+                .where(ArchivedMessage.room_id == room_pk)
+                .where(ArchivedMessage.legacy_id.isnot(None))
+            )
+            if source is not None:
+                q = q.where(ArchivedMessage.source == source)
+            return session.execute(
+                q.order_by(ArchivedMessage.timestamp.desc())
+            ).scalar()
+
+    def get_least_recent_with_legacy_id_after(
+        self, room_pk: int, after_id: str, source=ArchivedMessageSource.LIVE
+    ) -> ArchivedMessage | None:
+        with self.session() as session:
+            after_timestamp = (
+                session.query(ArchivedMessage.timestamp)
+                .filter(ArchivedMessage.legacy_id == after_id)
+                .scalar()
+            )
+            q = (
+                select(ArchivedMessage)
+                .where(ArchivedMessage.room_id == room_pk)
+                .where(ArchivedMessage.legacy_id.isnot(None))
+                .where(ArchivedMessage.source == source)
+                .where(ArchivedMessage.timestamp > after_timestamp)
+            )
+            return session.execute(q.order_by(ArchivedMessage.timestamp.asc())).scalar()
+
+    def get_by_legacy_id(self, room_pk: int, legacy_id: str) -> ArchivedMessage | None:
+        with self.session() as session:
+            return (
+                session.query(ArchivedMessage)
+                .filter(ArchivedMessage.room_id == room_pk)
+                .filter(ArchivedMessage.legacy_id == legacy_id)
+                .first()
+            )
+
 
 class MultiStore(EngineMixin):
     def get_xmpp_ids(self, user_pk: int, xmpp_id: str) -> list[str]:
@@ -676,27 +765,13 @@ class RoomStore(UpdatedMixin):
         super().__init__(*a, **kw)
         with self.session() as session:
             session.execute(
-                update(Room).values(subject_setter_id=None, user_resources=None)
+                update(Room).values(
+                    subject_setter=None, user_resources=None, history_filled=False
+                )
             )
             session.commit()
 
-    def add(self, user_pk: int, legacy_id: str, jid: JID) -> int:
-        if jid.resource:
-            raise TypeError
-        with self.session() as session:
-            existing = session.execute(
-                select(Room.id)
-                .where(Room.user_account_id == user_pk)
-                .where(Room.legacy_id == legacy_id)
-            ).scalar()
-            if existing is not None:
-                return existing
-            room = Room(jid=jid, user_account_id=user_pk, legacy_id=legacy_id)
-            session.add(room)
-            session.commit()
-            return room.id
-
-    def set_avatar(self, room_pk: int, avatar_pk: int) -> None:
+    def set_avatar(self, room_pk: int, avatar_pk: int | None) -> None:
         with self.session() as session:
             session.execute(
                 update(Room).where(Room.id == room_pk).values(avatar_id=avatar_pk)
@@ -728,43 +803,45 @@ class RoomStore(UpdatedMixin):
                 .where(Room.legacy_id == legacy_id)
             ).scalar()
 
-    def update(self, room: "LegacyMUC"):
-        from slidge.contact import LegacyContact
-
+    def update_subject_setter(self, room_pk: int, subject_setter: str | None):
         with self.session() as session:
-            if room.subject_setter is None:
-                subject_setter_id = None
-            elif isinstance(room.subject_setter, str):
-                subject_setter_id = None
-            elif isinstance(room.subject_setter, LegacyContact):
-                subject_setter_id = None
-            elif room.subject_setter.is_system:
-                subject_setter_id = None
-            else:
-                subject_setter_id = room.subject_setter.pk
-
             session.execute(
                 update(Room)
-                .where(Room.id == room.pk)
-                .values(
-                    updated=True,
-                    extra_attributes=room.serialize_extra_attributes(),
-                    name=room.name,
-                    description=room.description,
-                    user_resources=(
-                        None
-                        if not room._user_resources
-                        else json.dumps(list(room._user_resources))
-                    ),
-                    muc_type=room.type,
-                    subject=room.subject,
-                    subject_date=room.subject_date,
-                    subject_setter_id=subject_setter_id,
-                    participants_filled=room._participants_filled,
-                    n_participants=room._n_participants,
+                .where(Room.id == room_pk)
+                .values(subject_setter=subject_setter)
+            )
+            session.commit()
+
+    def update(self, muc: "LegacyMUC") -> int:
+        with self.session() as session:
+            if muc.pk is None:
+                row = Room(
+                    jid=muc.jid,
+                    legacy_id=str(muc.legacy_id),
+                    user_account_id=muc.user_pk,
                 )
+            else:
+                row = session.query(Room).filter(Room.id == muc.pk).one()
+
+            row.updated = True
+            row.extra_attributes = muc.serialize_extra_attributes()
+            row.name = muc.name
+            row.description = muc.description
+            row.user_resources = (
+                None
+                if not muc._user_resources
+                else json.dumps(list(muc._user_resources))
             )
+            row.muc_type = muc.type
+            row.subject = muc.subject
+            row.subject_date = muc.subject_date
+            row.subject_setter = muc.subject_setter
+            row.participants_filled = muc._participants_filled
+            row.n_participants = muc._n_participants
+            row.user_nick = muc.user_nick
+            session.add(row)
             session.commit()
+            return row.id
 
     def update_subject_date(
         self, room_pk: int, subject_date: Optional[datetime]
@@ -789,6 +866,11 @@ class RoomStore(UpdatedMixin):
             )
             session.commit()
 
+    def update_name(self, room_pk: int, name: Optional[str]) -> None:
+        with self.session() as session:
+            session.execute(update(Room).where(Room.id == room_pk).values(name=name))
+            session.commit()
+
     def update_n_participants(self, room_pk: int, n: Optional[int]) -> None:
         with self.session() as session:
             session.execute(
@@ -826,6 +908,20 @@ class RoomStore(UpdatedMixin):
                 is None
             )
 
+    def set_participants_filled(self, room_pk: int, val=True) -> None:
+        with self.session() as session:
+            session.execute(
+                update(Room).where(Room.id == room_pk).values(participants_filled=val)
+            )
+            session.commit()
+
+    def set_history_filled(self, room_pk: int, val=True) -> None:
+        with self.session() as session:
+            session.execute(
+                update(Room).where(Room.id == room_pk).values(history_filled=True)
+            )
+            session.commit()
+
     def get_all(self, user_pk: int) -> Iterator[Room]:
         with self.session() as session:
             yield from session.execute(
@@ -971,6 +1067,12 @@ class ParticipantStore(EngineMixin):
         with self.session() as session:
             session.execute(delete(Participant).where(Participant.id == participant_pk))
 
+    def get_count(self, room_pk: int) -> int:
+        with self.session() as session:
+            return session.query(
+                count(Participant.id).filter(Participant.room_id == room_pk)
+            ).scalar()
+
 
 log = logging.getLogger(__name__)
 _session: Optional[Session] = None

slidge/group/archive.py

@@ -6,8 +6,10 @@ from typing import TYPE_CHECKING, Collection, Optional
 
 from slixmpp import Iq, Message
 
+from ..db.models import ArchivedMessage, ArchivedMessageSource
 from ..db.store import MAMStore
 from ..util.archive_msg import HistoryMessage
+from ..util.types import HoleBound
 
 if TYPE_CHECKING:
     from .participant import LegacyParticipant
@@ -22,12 +24,16 @@ class MessageArchive:
         self,
         msg: Message,
         participant: Optional["LegacyParticipant"] = None,
+        archive_only=False,
+        legacy_msg_id=None,
     ):
         """
         Add a message to the archive if it is deemed archivable
 
         :param msg:
         :param participant:
+        :param archive_only:
+        :param legacy_msg_id:
         """
         if not archivable(msg):
             return
@@ -47,11 +53,50 @@ class MessageArchive:
                     "jid"
                 ] = f"{uuid.uuid4()}@{participant.xmpp.boundjid.bare}"
 
-        self.__store.add_message(self.room_pk, HistoryMessage(new_msg))
+        self.__store.add_message(
+            self.room_pk,
+            HistoryMessage(new_msg),
+            archive_only,
+            None if legacy_msg_id is None else str(legacy_msg_id),
+        )
 
     def __iter__(self):
         return iter(self.get_all())
 
+    @staticmethod
+    def __to_bound(stored: ArchivedMessage):
+        return HoleBound(
+            stored.legacy_id,  # type:ignore
+            stored.timestamp.replace(tzinfo=timezone.utc),
+        )
+
+    def get_hole_bounds(self) -> tuple[HoleBound | None, HoleBound | None]:
+        most_recent = self.__store.get_most_recent_with_legacy_id(self.room_pk)
+        if most_recent is None:
+            return None, None
+        if most_recent.source == ArchivedMessageSource.BACKFILL:
+            # most recent = only backfill, fetch everything since last backfill
+            return self.__to_bound(most_recent), None
+
+        most_recent_back_filled = self.__store.get_most_recent_with_legacy_id(
+            self.room_pk, ArchivedMessageSource.BACKFILL
+        )
+        if most_recent_back_filled is None:
+            # group was never back-filled, fetch everything before first live
+            least_recent_live = self.__store.get_first(self.room_pk, True)
+            assert least_recent_live is not None
+            return None, self.__to_bound(least_recent_live)
+
+        assert most_recent_back_filled.legacy_id is not None
+        least_recent_live = self.__store.get_least_recent_with_legacy_id_after(
+            self.room_pk, most_recent_back_filled.legacy_id
+        )
+        assert least_recent_live is not None
+        # this is a hole caused by slidge downtime
+        return self.__to_bound(most_recent_back_filled), self.__to_bound(
+            least_recent_live
+        )
+
     def get_all(
         self,
         start_date: Optional[datetime] = None,

slidge/group/bookmarks.py

@@ -3,6 +3,7 @@ import logging
 from typing import TYPE_CHECKING, Generic, Iterator, Optional, Type
 
 from slixmpp import JID
+from slixmpp.exceptions import XMPPError
 from slixmpp.jid import _unescape_node
 
 from ..contact.roster import ESCAPE_TABLE
@@ -57,15 +58,26 @@ class LegacyBookmarks(
         return f"<Bookmarks of {self.user_jid}>"
 
     async def __finish_init_muc(self, legacy_id: LegacyGroupIdType, jid: JID):
-        muc = self._muc_class(self.session, legacy_id=legacy_id, jid=jid)
         with self.__store.session():
-            muc.pk = self.__store.add(self.session.user_pk, str(muc.legacy_id), muc.jid)
-            muc.archive = MessageArchive(muc.pk, self.xmpp.store.mam)
-            await muc.avatar_wrap_update_info()
+            stored = self.__store.get_by_legacy_id(self.session.user_pk, str(legacy_id))
+            if stored is not None:
+                if stored.updated:
+                    return self._muc_class.from_store(self.session, stored)
+                muc = self._muc_class(self.session, legacy_id=legacy_id, jid=jid)
+                muc.pk = stored.id
+            else:
+                muc = self._muc_class(self.session, legacy_id=legacy_id, jid=jid)
+
+            try:
+                with muc.updating_info():
+                    await muc.avatar_wrap_update_info()
+            except Exception as e:
+                raise XMPPError("internal-server-error", str(e))
             if not muc.user_nick:
                 muc.user_nick = self._user_nick
             self.log.debug("MUC created: %r", muc)
-            self.__store.update(muc)
+            muc.pk = self.__store.update(muc)
+            muc.archive = MessageArchive(muc.pk, self.xmpp.store.mam)
         return muc
 
     async def legacy_id_to_jid_local_part(self, legacy_id: LegacyGroupIdType):

slidge/group/participant.py

@@ -117,6 +117,8 @@ class LegacyParticipant(
         if self._affiliation == affiliation:
             return
         self._affiliation = affiliation
+        if not self.muc._participants_filled:
+            return
         self.__part_store.set_affiliation(self.pk, affiliation)
         if not self._presence_sent:
             return
@@ -145,6 +147,8 @@ class LegacyParticipant(
         if self._role == role:
             return
         self._role = role
+        if not self.muc._participants_filled:
+            return
         self.__part_store.set_role(self.pk, role)
         if not self._presence_sent:
             return
@@ -154,6 +158,8 @@ class LegacyParticipant(
         if self._hats == hats:
             return
         self._hats = hats
+        if not self.muc._participants_filled:
+            return
         self.__part_store.set_hats(self.pk, hats)
         if not self._presence_sent:
             return
@@ -313,6 +319,7 @@ class LegacyParticipant(
         stanza: MessageOrPresenceTypeVar,
         full_jid: Optional[JID] = None,
         archive_only=False,
+        legacy_msg_id=None,
         **send_kwargs,
     ) -> MessageOrPresenceTypeVar:
         stanza["occupant-id"]["id"] = self.__occupant_id
@@ -331,8 +338,8 @@ class LegacyParticipant(
             else:
                 stanza.send()
         else:
-            if isinstance(stanza, Message):
-                self.muc.archive.add(stanza, self)
+            if hasattr(self.muc, "archive") and isinstance(stanza, Message):
+                self.muc.archive.add(stanza, self, archive_only, legacy_msg_id)
             if archive_only:
                 return stanza
             for user_full_jid in self.muc.user_full_jids():
@@ -460,7 +467,7 @@ class LegacyParticipant(
     ):
         if update_muc:
             self.muc._subject = subject  # type: ignore
-            self.muc.subject_setter = self
+            self.muc.subject_setter = self.nickname
             self.muc.subject_date = when
 
         msg = self._make_message()