Coverage for slidge/core/collection.py: 93%
140 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-10 12:26 +0000
1import abc
2from collections.abc import AsyncIterator, Iterator
3from contextlib import contextmanager
4from typing import Any, Literal, Protocol
6import sqlalchemy as sa
7import sqlalchemy.orm
8from slixmpp import JID
9from slixmpp.exceptions import XMPPError
11from slidge.db.meta import Base
13from ..contact import LegacyContact
14from ..db.models import Contact, GatewayUser, Room, Space
15from ..util.lock import NamedLockMixin
16from ..util.types import AnyRoom, AnySession, AnySpace
19class _Store[M](Protocol):
20 def by_legacy_id(
21 self, orm: sa.orm.Session, user_pk: int, legacy_id: str
22 ) -> M | None: ...
25class _ChildFactory[M, C](Protocol):
26 @classmethod
27 @contextmanager
28 def _batch_set_attributes(cls) -> Iterator[None]: ...
30 def __call__(self, session: AnySession, *, stored: M) -> C: ...
33class _Child(Protocol):
34 @property
35 def _stored(self) -> Base: ...
37 def _on_delete(self) -> None: ...
40class BaseByLegacyID[
41 ModelT: Contact | Room | Space,
42 ChildT: LegacyContact | AnyRoom | AnySpace,
43](abc.ABC, NamedLockMixin):
44 def __init__(
45 self,
46 session: AnySession,
47 model: type[ModelT],
48 child_cls: _ChildFactory[ModelT, ChildT],
49 store: _Store[ModelT],
50 ) -> None:
51 super().__init__()
52 self._model = model
53 self.__child_cls = child_cls
54 self._store = store
55 self._session = session
56 self.xmpp = session.xmpp
57 self.__fill_cache_db: dict[str, ModelT] = {}
59 # Any here to prevent mypy from inferring invariance, we need covariance
60 # to work with LegacyRoster
61 self.__fill_cache_children: dict[str, Any] = {}
62 self.__filling: bool = False
64 @property
65 def _user(self) -> GatewayUser:
66 return self._session.user
68 def _orm(self, expire_on_commit: bool = False) -> sa.orm.Session:
69 return self.xmpp.store.session(expire_on_commit=expire_on_commit)
71 def _from_store(self, stored: ModelT) -> ChildT:
72 return self.__child_cls(self._session, stored=stored)
74 def __repr__(self) -> str:
75 return f"<{self.__class__} of {self._session.user_jid}>"
77 def __iter__(self) -> Iterator[ChildT]:
78 with self._orm() as orm:
79 for stored in orm.query(self._model).filter_by(user=self._user).all():
80 yield self._from_store(stored)
82 @property
83 def _db_rows(self) -> Iterator[ModelT]:
84 with self._orm() as orm:
85 yield from orm.query(self._model).filter_by(user=self._user).all()
87 def get(self, legacy_id: str) -> ChildT | None:
88 """
89 Get a :term:`Contact` or :term:`Room` that has been created before.
90 """
91 if self.__filling:
92 child = self.__fill_cache_children.get(legacy_id)
93 if child is not None:
94 return child # type:ignore[no-any-return]
96 stored = self.__fill_cache_db.get(legacy_id)
97 if stored is None:
98 return None
100 self.__fill_cache_children[legacy_id] = self._from_store(stored)
101 return self.__fill_cache_children[legacy_id] # type:ignore[no-any-return]
103 with self._orm() as orm:
104 stored = self._store.by_legacy_id(orm, self._session.user_pk, legacy_id)
105 return None if stored is None else self._from_store(stored)
107 def _create(self, legacy_id: str, **kwargs: object) -> ChildT:
108 if not legacy_id:
109 raise XMPPError(
110 "internal-server-error",
111 f"Attempted to create a {self._model.__class__} with legacy_id=''",
112 )
113 stored = self._model(legacy_id=legacy_id, user=self._user, **kwargs)
114 return self._commit(stored, legacy_id)
116 def _commit(self, stored: ModelT, legacy_id: str) -> ChildT:
117 with self._orm() as orm:
118 orm.add(stored)
119 orm.commit()
120 child = self._from_store(stored)
121 if self.is_filling:
122 self.__fill_cache_db[legacy_id] = stored
123 self.__fill_cache_children[legacy_id] = child
124 return child
126 async def create(self, legacy_id: str) -> ChildT:
127 """
128 Create a new :term:`Contact`, :term:`Space` or :term:`Space`.
130 This must raise if legacy id does not map to something existing.
132 A typical pattern is to have a method constructing a ``Contact`` or
133 a ``Room`` from a "metadata struct" provided by the legacy network,
134 in the spirit of:
136 .. code-block:: python3
138 class LegacyRoster[Contact]:
139 def create(self, legacy_id: str) -> Contact:
140 meta = self._session.legacy_client.fetch_contact(legacy_id)
141 if meta is None:
142 raise XMPPError("item-not-found", f"No contact {legacy_id!r}")
143 return await self.get_or_create_by_meta(meta)
145 def get_or_create_by_meta(self, meta: legacy_lib.ContactMetadata) -> Contact:
146 contact = self.get(meta.id) or self._create(legacy_id)
147 contact.set_meta(meta)
148 return contact
149 """
150 return self._create(legacy_id)
152 async def get_or_create(self, legacy_id: str, *args: object) -> ChildT:
153 """
154 Get an existing :term:`Recipient` or :term:`Space` or :meth:`.create`
155 it if does not exist.
157 :param legacy_id: The unique ID of this Contact
158 :param *args: Optionally, arguments passed to :meth:`.create`. Make
159 sure that these match whatever your :meth:`.create` override
160 expects. NB: nothing is passed to :meth:`.get`.
161 """
162 async with self.lock(legacy_id):
163 return self.get(legacy_id) or await self.create(legacy_id, *args)
165 def delete(self, /, child: _Child) -> None:
166 child._on_delete()
167 with self._orm() as orm:
168 orm.delete(child._stored)
169 orm.commit()
171 async def _batch_fill(self) -> AsyncIterator[ChildT]:
172 updates = []
173 with self._preload():
174 async for child in self.fill():
175 yield child
176 update = child._pending_updates
177 if update:
178 updates.append({"id": child._stored.id, **update})
179 with self._orm() as orm:
180 orm.execute(sa.update(self._model), updates)
181 orm.commit()
183 @contextmanager
184 def _preload(self) -> Iterator[None]:
185 self.__filling = True
186 try:
187 with self._orm() as orm:
188 self.__fill_cache_db = {
189 r.legacy_id: r for r in orm.scalars(self._batch_load())
190 }
191 with self.__child_cls._batch_set_attributes():
192 yield
193 finally:
194 self.__filling = False
195 self.__fill_cache_db.clear()
196 self.__fill_cache_children.clear()
198 def _batch_load(self) -> sa.Select[ModelT]:
199 return (
200 sa.select(self._model)
201 .where(self._model.user == self._user)
202 .options(
203 sa.orm.joinedload(self._model.avatar),
204 )
205 )
207 async def fill(self) -> AsyncIterator[ChildT]:
208 """
209 This is called once on startup, and should yield :term:`recipients <Recipient>`
210 with updated metadata.
211 """
212 return
213 yield
215 @property
216 def is_filling(self) -> bool:
217 return self.__filling
220class BaseLocalPart[
221 ModelT: Contact | Room,
222 ChildT: LegacyContact | AnyRoom,
223](BaseByLegacyID[ModelT, ChildT]):
224 def __init__(
225 self,
226 session: AnySession,
227 model: type[ModelT],
228 child_cls: _ChildFactory[ModelT, ChildT],
229 store: _Store[ModelT],
230 legacy_type: Literal["contact", "room"],
231 ) -> None:
232 super().__init__(session, model, child_cls, store)
233 self._legacy_type = legacy_type
235 def _create(self, legacy_id: str, **kwargs: object) -> ChildT:
236 stored = self._model(
237 legacy_id=legacy_id,
238 user=self._user,
239 jid_localpart=self.xmpp.legacy_id_to_localpart(
240 self._legacy_type, legacy_id
241 ),
242 **kwargs,
243 )
244 return self._commit(stored, legacy_id)
246 async def get_or_create_by_jid(self, jid: JID | str) -> ChildT:
247 jid = JID(JID(jid).bare)
248 localpart = jid.local
249 if not localpart:
250 raise XMPPError(
251 "bad-request",
252 f"A {self._legacy_type} must have a local part in their JID",
253 )
254 type_, legacy_id = self.xmpp.localpart_to_legacy_id(localpart)
255 if type_ != self._legacy_type:
256 raise XMPPError(
257 "bad-request",
258 f"Expected JID for {self._legacy_type}, not a {type_} JID",
259 )
260 return await self.get_or_create(legacy_id)
262 def get_by_jid(self, jid: JID) -> ChildT | None:
263 with self._orm() as orm:
264 stored = (
265 orm.query(self._model)
266 .filter_by(user=self._user, jid_localpart=jid.local)
267 .one_or_none()
268 )
269 if stored is not None:
270 return self._from_store(stored)
271 return None