slidge 0.2.12__py3-none-any.whl → 0.3.0__py3-none-any.whl

slidge/core/mixins/attachment.py

@@ -11,7 +11,7 @@ from datetime import datetime
 from itertools import chain
 from mimetypes import guess_extension, guess_type
 from pathlib import Path
-from typing import IO, AsyncIterator, Collection, Optional, Sequence, Union
+from typing import Collection, Optional, Sequence, Union
 from urllib.parse import quote as urlquote
 from uuid import uuid4
 from xml.etree import ElementTree as ET
@@ -25,6 +25,7 @@ from slixmpp.plugins.xep_0363 import FileUploadError
 from slixmpp.plugins.xep_0447.stanza import StatelessFileSharing
 
 from ...db.avatar import avatar_cache
+from ...db.models import Attachment
 from ...util.types import (
     LegacyAttachment,
     LegacyMessageType,
@@ -37,9 +38,7 @@ from .message_text import TextMessageMixin
 
 
 class AttachmentMixin(TextMessageMixin):
-    def __init__(self, *a, **kw):
-        super().__init__(*a, **kw)
-        self.__store = self.xmpp.store.attachments
+    is_group: bool
 
     async def __upload(
         self,
@@ -137,28 +136,44 @@ class AttachmentMixin(TextMessageMixin):
 
         return destination, uploaded_url
 
+    async def __valid_url(self, url: str) -> bool:
+        async with self.session.http.head(url) as r:
+            return r.status < 400
+
+    async def __get_stored(self, attachment: LegacyAttachment) -> Attachment:
+        if attachment.legacy_file_id is not None and self.session is not NotImplemented:
+            with self.xmpp.store.session() as orm:
+                stored = (
+                    orm.query(Attachment)
+                    .filter_by(
+                        legacy_file_id=str(attachment.legacy_file_id),
+                        user_account_id=self.session.user_pk,
+                    )
+                    .one_or_none()
+                )
+                if stored is not None:
+                    if not await self.__valid_url(stored.url):
+                        stored.url = None  # type:ignore
+                    return stored
+        return Attachment(
+            user_account_id=None
+            if self.session is NotImplemented
+            else self.session.user_pk,
+            legacy_file_id=None
+            if attachment.legacy_file_id is None
+            else str(attachment.legacy_file_id),
+            url=attachment.url,
+        )
+
     async def __get_url(
-        self,
-        file_path: Optional[Path] = None,
-        async_data_stream: Optional[AsyncIterator[bytes]] = None,
-        data_stream: Optional[IO[bytes]] = None,
-        data: Optional[bytes] = None,
-        file_url: Optional[str] = None,
-        file_name: Optional[str] = None,
-        content_type: Optional[str] = None,
-        legacy_file_id: Optional[Union[str, int]] = None,
+        self, attachment: LegacyAttachment, stored: Attachment
     ) -> tuple[bool, Optional[Path], str]:
-        if legacy_file_id:
-            cache = self.__store.get_url(str(legacy_file_id))
-            if cache is not None:
-                async with self.session.http.head(cache) as r:
-                    if r.status < 400:
-                        return False, None, cache
-                    else:
-                        self.__store.remove(str(legacy_file_id))
-
-        if file_url and config.USE_ATTACHMENT_ORIGINAL_URLS:
-            return False, None, file_url
+        if attachment.url and config.USE_ATTACHMENT_ORIGINAL_URLS:
+            return False, None, attachment.url
+
+        file_name = attachment.name
+        content_type = attachment.content_type
+        file_path = attachment.path
 
         if file_name and len(file_name) > config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH:
             log.debug("Trimming long filename: %s", file_name)
@@ -176,47 +191,48 @@ class AttachmentMixin(TextMessageMixin):
                         file_name += ext
             temp_dir = Path(tempfile.mkdtemp())
             file_path = temp_dir / file_name
-            if file_url:
-                async with self.session.http.get(file_url) as r:
+            if attachment.url:
+                async with self.session.http.get(attachment.url) as r:
                     r.raise_for_status()
                     with file_path.open("wb") as f:
                         f.write(await r.read())
 
-            elif data_stream is not None:
-                data = data_stream.read()
+            elif attachment.stream is not None:
+                data = attachment.stream.read()
                 if data is None:
                     raise RuntimeError
 
                 with file_path.open("wb") as f:
                     f.write(data)
-            elif async_data_stream is not None:
+            elif attachment.aio_stream is not None:
                 # TODO: patch slixmpp to allow this as data source for
                 #       upload_file() so we don't even have to write anything
                 #       to disk.
                 with file_path.open("wb") as f:
-                    async for chunk in async_data_stream:
+                    async for chunk in attachment.aio_stream:
                         f.write(chunk)
-            elif data is not None:
+            elif attachment.data is not None:
                 with file_path.open("wb") as f:
-                    f.write(data)
+                    f.write(attachment.data)
 
             is_temp = not bool(config.NO_UPLOAD_PATH)
         else:
             is_temp = False
 
+        assert isinstance(file_path, Path)
         if config.FIX_FILENAME_SUFFIX_MIME_TYPE:
             file_name = str(fix_suffix(file_path, content_type, file_name))
 
         if config.NO_UPLOAD_PATH:
             local_path, new_url = await self.__no_upload(
-                file_path, file_name, legacy_file_id
+                file_path, file_name, stored.legacy_file_id
             )
             new_url = (config.NO_UPLOAD_URL_PREFIX or "") + "/" + urlquote(new_url)
         else:
             local_path = file_path
             new_url = await self.__upload(file_path, file_name, content_type)
-        if legacy_file_id and new_url is not None:
-            self.__store.set_url(self.session.user_pk, str(legacy_file_id), new_url)
+        if stored.legacy_file_id and new_url is not None:
+            stored.url = new_url
 
         return is_temp, local_path, new_url
 
@@ -225,13 +241,11 @@ class AttachmentMixin(TextMessageMixin):
         msg: Message,
         uploaded_url: str,
         path: Optional[Path],
-        content_type: Optional[str] = None,
-        caption: Optional[str] = None,
-        file_name: Optional[str] = None,
+        attachment: LegacyAttachment,
+        stored: Attachment,
     ) -> Thumbnail | None:
-        cache = self.__store.get_sims(uploaded_url)
-        if cache:
-            ref = self.xmpp["xep_0372"].stanza.Reference(xml=ET.fromstring(cache))
+        if stored.sims is not None:
+            ref = self.xmpp["xep_0372"].stanza.Reference(xml=ET.fromstring(stored.sims))
             msg.append(ref)
             if ref["sims"]["file"].get_plugin("thumbnail", check=True):
                 return ref["sims"]["file"]["thumbnail"]
@@ -242,12 +256,14 @@ class AttachmentMixin(TextMessageMixin):
             return None
 
         ref = self.xmpp["xep_0385"].get_sims(
-            path, [uploaded_url], content_type, caption
+            path, [uploaded_url], attachment.content_type, attachment.caption
         )
-        if file_name:
-            ref["sims"]["file"]["name"] = file_name
+        if attachment.name:
+            ref["sims"]["file"]["name"] = attachment.name
         thumbnail = None
-        if content_type is not None and content_type.startswith("image"):
+        if attachment.content_type is not None and attachment.content_type.startswith(
+            "image"
+        ):
             try:
                 h, x, y = await self.xmpp.loop.run_in_executor(
                     avatar_cache._thread_pool, get_thumbhash, path
@@ -261,8 +277,7 @@ class AttachmentMixin(TextMessageMixin):
                 thumbnail["media-type"] = "image/thumbhash"
                 thumbnail["uri"] = "data:image/thumbhash;base64," + urlquote(h)
 
-        self.__store.set_sims(uploaded_url, str(ref))
-
+        stored.sims = str(ref)
         msg.append(ref)
 
         return thumbnail
@@ -272,26 +287,29 @@ class AttachmentMixin(TextMessageMixin):
         msg: Message,
         uploaded_url: str,
         path: Optional[Path],
-        content_type: Optional[str] = None,
-        caption: Optional[str] = None,
-        file_name: Optional[str] = None,
+        attachment: LegacyAttachment,
+        stored: Attachment,
         thumbnail: Optional[Thumbnail] = None,
-    ):
-        cache = self.__store.get_sfs(uploaded_url)
-        if cache:
-            msg.append(StatelessFileSharing(xml=ET.fromstring(cache)))
+    ) -> None:
+        if stored.sfs is not None:
+            msg.append(StatelessFileSharing(xml=ET.fromstring(stored.sfs)))
             return
 
         if not path:
             return
 
-        sfs = self.xmpp["xep_0447"].get_sfs(path, [uploaded_url], content_type, caption)
-        if file_name:
-            sfs["file"]["name"] = file_name
+        sfs = self.xmpp["xep_0447"].get_sfs(
+            path, [uploaded_url], attachment.content_type, attachment.caption
+        )
+        if attachment.name:
+            sfs["file"]["name"] = attachment.name
+        if attachment.disposition:
+            sfs["disposition"] = attachment.disposition
+        else:
+            del sfs["disposition"]
         if thumbnail is not None:
             sfs["file"].append(thumbnail)
-        self.__store.set_sfs(uploaded_url, str(sfs))
-
+        stored.sfs = str(sfs)
         msg.append(sfs)
 
     def __send_url(
@@ -300,9 +318,9 @@ class AttachmentMixin(TextMessageMixin):
         legacy_msg_id: LegacyMessageType,
         uploaded_url: str,
         caption: Optional[str] = None,
-        carbon=False,
+        carbon: bool = False,
         when: Optional[datetime] = None,
-        correction=False,
+        correction: bool = False,
         **kwargs,
     ) -> list[Message]:
         msg["oob"]["url"] = uploaded_url
@@ -325,51 +343,21 @@ class AttachmentMixin(TextMessageMixin):
                 self._set_msg_id(msg, legacy_msg_id)
             return [self._send(msg, carbon=carbon, **kwargs)]
 
-    async def send_file(
+    def __get_base_message(
         self,
-        file_path: Optional[Union[Path, str]] = None,
         legacy_msg_id: Optional[LegacyMessageType] = None,
-        *,
-        async_data_stream: Optional[AsyncIterator[bytes]] = None,
-        data_stream: Optional[IO[bytes]] = None,
-        data: Optional[bytes] = None,
-        file_url: Optional[str] = None,
-        file_name: Optional[str] = None,
-        content_type: Optional[str] = None,
         reply_to: Optional[MessageReference] = None,
         when: Optional[datetime] = None,
-        caption: Optional[str] = None,
-        legacy_file_id: Optional[Union[str, int]] = None,
         thread: Optional[LegacyThreadType] = None,
-        **kwargs,
-    ) -> tuple[Optional[str], list[Message]]:
-        """
-        Send a single file from this :term:`XMPP Entity`.
-
-        :param file_path: Path to the attachment
-        :param async_data_stream: Alternatively (and ideally) an AsyncIterator yielding bytes
-        :param data_stream: Alternatively, a stream of bytes (such as a File object)
-        :param data: Alternatively, a bytes object
-        :param file_url: Alternatively, a URL
-        :param file_name: How the file should be named.
-        :param content_type: MIME type, inferred from filename if not given
-        :param legacy_msg_id: If you want to be able to transport read markers from the gateway
-            user to the legacy network, specify this
-        :param reply_to: Quote another message (:xep:`0461`)
-        :param when: when the file was sent, for a "delay" tag (:xep:`0203`)
-        :param caption: an optional text that is linked to the file
-        :param legacy_file_id: A unique identifier for the file on the legacy network.
-             Plugins should try their best to provide it, to avoid duplicates.
-        :param thread:
-        """
-        carbon = kwargs.pop("carbon", False)
-        mto = kwargs.pop("mto", None)
-        store_multi = kwargs.pop("store_multi", True)
-        correction = kwargs.get("correction", False)
+        carbon: bool = False,
+        correction: bool = False,
+        mto: Optional[JID] = None,
+    ) -> Message:
         if correction and (original_xmpp_id := self._legacy_to_xmpp(legacy_msg_id)):
-            xmpp_ids = self.xmpp.store.multi.get_xmpp_ids(
-                self.session.user_pk, original_xmpp_id
-            )
+            with self.xmpp.store.session() as orm:
+                xmpp_ids = self.xmpp.store.id_map.get_xmpp(
+                    orm, self._recipient_pk(), str(legacy_msg_id), self.is_group
+                )
 
             for xmpp_id in xmpp_ids:
                 if xmpp_id == original_xmpp_id:
@@ -386,7 +374,7 @@ class AttachmentMixin(TextMessageMixin):
         else:
             reply_to_for_attachment = reply_to
 
-        msg = self._make_message(
+        return self._make_message(
             when=when,
             reply_to=reply_to_for_attachment,
             carbon=carbon,
@@ -394,43 +382,116 @@ class AttachmentMixin(TextMessageMixin):
             thread=thread,
         )
 
-        if content_type is None and (name := (file_name or file_path or file_url)):
-            content_type, _ = guess_type(name)
-
-        is_temp, local_path, new_url = await self.__get_url(
-            Path(file_path) if file_path else None,
-            async_data_stream,
-            data_stream,
-            data,
-            file_url,
-            file_name,
-            content_type,
-            legacy_file_id,
+    async def send_file(
+        self,
+        attachment: LegacyAttachment | Path | str,
+        legacy_msg_id: Optional[LegacyMessageType] = None,
+        *,
+        reply_to: Optional[MessageReference] = None,
+        when: Optional[datetime] = None,
+        thread: Optional[LegacyThreadType] = None,
+        **kwargs,
+    ) -> tuple[Optional[str], list[Message]]:
+        """
+        Send a single file from this :term:`XMPP Entity`.
+
+        :param attachment: The file to send.
+            Ideally, a :class:`.LegacyAttachment` with a unique ``legacy_file_id``
+            attribute set, to optimise potential future reuses.
+            It can also be:
+            - a :class:`pathlib.Path` instance to point to a local file, or
+            - a ``str``, representing a fetchable HTTP URL.
+        :param legacy_msg_id: If you want to be able to transport read markers from the gateway
+            user to the legacy network, specify this
+        :param reply_to: Quote another message (:xep:`0461`)
+        :param when: when the file was sent, for a "delay" tag (:xep:`0203`)
+        :param thread:
+        """
+        coro = self.__send_file(
+            attachment,
+            legacy_msg_id,
+            reply_to=reply_to,
+            when=when,
+            thread=thread,
+            **kwargs,
         )
+        if self.session is NotImplemented:
+            return await coro
+        elif not isinstance(attachment, LegacyAttachment):
+            return await coro
+        elif attachment.legacy_file_id is None:
+            return await coro
+        else:
+            # prevents race conditions where we download the same thing several time
+            # and end up attempting to insert it twice in the DB, raising an
+            # IntegrityError.
+            async with self.session.lock(("attachment", attachment.legacy_file_id)):
+                return await coro
 
-        if new_url is None:
-            msg["body"] = (
-                "I tried to send a file, but something went wrong. "
-                "Tell your slidge admin to check the logs."
-            )
-            self._set_msg_id(msg, legacy_msg_id)
-            return None, [self._send(msg, **kwargs)]
+    async def __send_file(
+        self,
+        attachment: LegacyAttachment | Path | str,
+        legacy_msg_id: Optional[LegacyMessageType] = None,
+        *,
+        reply_to: Optional[MessageReference] = None,
+        when: Optional[datetime] = None,
+        thread: Optional[LegacyThreadType] = None,
+        **kwargs,
+    ) -> tuple[Optional[str], list[Message]]:
+        store_multi = kwargs.pop("store_multi", True)
+        carbon = kwargs.pop("carbon", False)
+        mto = kwargs.pop("mto", None)
+        correction = kwargs.get("correction", False)
 
-        thumbnail = await self.__set_sims(
-            msg, new_url, local_path, content_type, caption, file_name
-        )
-        self.__set_sfs(
-            msg, new_url, local_path, content_type, caption, file_name, thumbnail
+        msg = self.__get_base_message(
+            legacy_msg_id, reply_to, when, thread, carbon, correction, mto
         )
+
+        if isinstance(attachment, str):
+            attachment = LegacyAttachment(url=attachment)
+        elif isinstance(attachment, Path):
+            attachment = LegacyAttachment(path=attachment)
+
+        stored = await self.__get_stored(attachment)
+
+        if attachment.content_type is None and (
+            name := (attachment.name or attachment.url or attachment.path)
+        ):
+            attachment.content_type, _ = guess_type(name)
+
+        if stored.url:
+            is_temp = False
+            local_path = None
+            new_url = stored.url
+        else:
+            is_temp, local_path, new_url = await self.__get_url(attachment, stored)
+            if new_url is None:
+                msg["body"] = (
+                    "I tried to send a file, but something went wrong. "
+                    "Tell your slidge admin to check the logs."
+                )
+                self._set_msg_id(msg, legacy_msg_id)
+                return None, [self._send(msg, **kwargs)]
+
+        stored.url = new_url
+        thumbnail = await self.__set_sims(msg, new_url, local_path, attachment, stored)
+        self.__set_sfs(msg, new_url, local_path, attachment, stored, thumbnail)
+
+        if self.session is not NotImplemented:
+            with self.xmpp.store.session(expire_on_commit=False) as orm:
+                orm.add(stored)
+                orm.commit()
+
         if is_temp and isinstance(local_path, Path):
             local_path.unlink()
             local_path.parent.rmdir()
 
         msgs = self.__send_url(
-            msg, legacy_msg_id, new_url, caption, carbon, when, **kwargs
+            msg, legacy_msg_id, new_url, attachment.caption, carbon, when, **kwargs
         )
-        if store_multi:
-            self.__store_multi(legacy_msg_id, msgs)
+        if self.session is not NotImplemented:
+            if store_multi:
+                self.__store_multi(legacy_msg_id, msgs)
         return new_url, msgs
 
     def __send_body(
@@ -463,18 +524,15 @@ class AttachmentMixin(TextMessageMixin):
         reply_to: Optional[MessageReference] = None,
         when: Optional[datetime] = None,
         thread: Optional[LegacyThreadType] = None,
-        body_first=False,
-        correction=False,
+        body_first: bool = False,
+        correction: bool = False,
         correction_event_id: Optional[LegacyMessageType] = None,
         **kwargs,
-    ):
+    ) -> None:
         # TODO: once the epic XEP-0385 vs XEP-0447 battle is over, pick
         #       one and stop sending several attachments this way
         # we attach the legacy_message ID to the last message we send, because
         # we don't want several messages with the same ID (especially for MUC MAM)
-        # TODO: refactor this so we limit the number of SQL calls, ie, if
-        #       the legacy file ID is known, only fetch the row once, and if it
-        #       is new, write it all in a single call
         if not attachments and not body:
             # ignoring empty message
             return
@@ -500,18 +558,11 @@ class AttachmentMixin(TextMessageMixin):
             else:
                 legacy = None
             _url, msgs = await self.send_file(
-                file_path=attachment.path,
-                legacy_msg_id=legacy,
-                file_url=attachment.url,
-                data_stream=attachment.stream,
-                data=attachment.data,
+                attachment,
+                legacy,
                 reply_to=reply_to,
                 when=when,
                 thread=thread,
-                file_name=attachment.name,
-                content_type=attachment.content_type,
-                legacy_file_id=attachment.legacy_file_id,
-                caption=attachment.caption,
                 store_multi=False,
                 **kwargs,
             )
@@ -524,7 +575,7 @@ class AttachmentMixin(TextMessageMixin):
         self,
         legacy_msg_id: Optional[LegacyMessageType],
         all_msgs: Sequence[Optional[Message]],
-    ):
+    ) -> None:
         if legacy_msg_id is None:
             return
         ids = []
@@ -535,9 +586,11 @@ class AttachmentMixin(TextMessageMixin):
                 ids.append(stanza_id["id"])
             else:
                 ids.append(msg.get_id())
-        self.xmpp.store.multi.set_xmpp_ids(
-            self.session.user_pk, str(legacy_msg_id), ids
-        )
+        with self.xmpp.store.session() as orm:
+            self.xmpp.store.id_map.set_msg(
+                orm, self._recipient_pk(), str(legacy_msg_id), ids, self.is_group
+            )
+            orm.commit()
 
 
 def get_thumbhash(path: Path) -> tuple[str, int, int]: