From 601058d61cbd514478d671287e563b7b0ead1403 Mon Sep 17 00:00:00 2001 From: Sumner Evans Date: Fri, 25 Mar 2022 14:49:45 -0600 Subject: [PATCH] async media: add ability to upload media asynchronously Co-authored-by: Tulir Asokan --- mautrix_telegram/example-config.yaml | 3 +++ mautrix_telegram/portal.py | 20 ++++++++++++++++--- mautrix_telegram/puppet.py | 1 + mautrix_telegram/util/file_transfer.py | 19 ++++++++++++++---- .../util/parallel_file_transfer.py | 7 ++++++- requirements.txt | 2 +- 6 files changed, 43 insertions(+), 9 deletions(-) diff --git a/mautrix_telegram/example-config.yaml b/mautrix_telegram/example-config.yaml index 29552262..e1f70a79 100644 --- a/mautrix_telegram/example-config.yaml +++ b/mautrix_telegram/example-config.yaml @@ -16,6 +16,9 @@ homeserver: status_endpoint: null # Endpoint for reporting per-message status. message_send_checkpoint_endpoint: null + # Whether asynchronous uploads via MSC2246 should be enabled for media. + # Requires a media repo that supports MSC2246. + async_media: false # Application service host/registration related details # Changing these values requires regeneration of the registration. diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index a50b57c2..7a32e7ad 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -1212,7 +1212,12 @@ class Portal(DBPortal, BasePortal): self.photo_id = "" self.avatar_url = None elif self.photo_id != photo_id or not self.avatar_url: - file = await util.transfer_file_to_matrix(user.client, self.main_intent, loc) + file = await util.transfer_file_to_matrix( + user.client, + self.main_intent, + loc, + async_upload=self.config["homeserver.async_media"], + ) if not file: return False self.photo_id = photo_id @@ -2224,7 +2229,11 @@ class Portal(DBPortal, BasePortal): ) return await self._send_message(intent, content, timestamp=evt.date) file = await util.transfer_file_to_matrix( - source.client, intent, loc, encrypt=self.encrypted + source.client, + intent, + loc, + encrypt=self.encrypted, + async_upload=self.config["homeserver.async_media"], ) if not file: return None @@ -2390,6 +2399,7 @@ class Portal(DBPortal, BasePortal): filename=attrs.name, parallel_id=parallel_id, encrypt=self.encrypted, + async_upload=self.config["homeserver.async_media"], ) if not file: return None @@ -2527,7 +2537,11 @@ class Portal(DBPortal, BasePortal): beeper_link_preview["og:image:height"] = largest_size.h beeper_link_preview["og:image:width"] = largest_size.w file = await util.transfer_file_to_matrix( - source.client, intent, loc, encrypt=self.encrypted + source.client, + intent, + loc, + encrypt=self.encrypted, + async_upload=self.config["homeserver.async_media"], ) if file.decryption_info: diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index b65ab897..a3e3760d 100644 --- a/mautrix_telegram/puppet.py +++ b/mautrix_telegram/puppet.py @@ -359,6 +359,7 @@ class Puppet(DBPuppet, BasePuppet): location=InputPeerPhotoFileLocation( peer=await self.get_input_entity(source), photo_id=photo.photo_id, big=True ), + async_upload=self.config["homeserver.async_media"], ) if not file: return False diff --git a/mautrix_telegram/util/file_transfer.py b/mautrix_telegram/util/file_transfer.py index 3a48e95c..e9bcb32f 100644 --- a/mautrix_telegram/util/file_transfer.py +++ b/mautrix_telegram/util/file_transfer.py @@ -144,6 +144,7 @@ async def transfer_thumbnail_to_matrix( custom_data: bytes | None = None, width: int | None = None, height: int | None = None, + async_upload: bool = False, ) -> DBTelegramFile | None: if not Image or not VideoFileClip: return None @@ -178,7 +179,7 @@ async def transfer_thumbnail_to_matrix( if encrypt: file, decryption_info = encrypt_attachment(file) upload_mime_type = "application/octet-stream" - content_uri = await intent.upload_media(file, upload_mime_type) + content_uri = await intent.upload_media(file, upload_mime_type, async_upload=async_upload) if decryption_info: decryption_info.url = content_uri @@ -220,6 +221,7 @@ async def transfer_file_to_matrix( filename: str | None = None, encrypt: bool = False, parallel_id: int | None = None, + async_upload: bool = False, ) -> DBTelegramFile | None: location_id = _location_to_id(location) if not location_id: @@ -246,6 +248,7 @@ async def transfer_file_to_matrix( filename, encrypt, parallel_id, + async_upload=async_upload, ) @@ -260,6 +263,7 @@ async def _unlocked_transfer_file_to_matrix( filename: str | None, encrypt: bool, parallel_id: int | None, + async_upload: bool = False, ) -> DBTelegramFile | None: db_file = await DBTelegramFile.get(loc_id) if db_file: @@ -269,7 +273,7 @@ async def _unlocked_transfer_file_to_matrix( if parallel_id and isinstance(location, Document) and (not is_sticker or not tgs_convert): db_file = await parallel_transfer_to_matrix( - client, intent, loc_id, location, filename, encrypt, parallel_id + client, intent, loc_id, location, filename, encrypt, parallel_id, async_upload ) mime_type = location.mime_type file = None @@ -305,7 +309,7 @@ async def _unlocked_transfer_file_to_matrix( if encrypt and encrypt_attachment: file, decryption_info = encrypt_attachment(file) upload_mime_type = "application/octet-stream" - content_uri = await intent.upload_media(file, upload_mime_type) + content_uri = await intent.upload_media(file, upload_mime_type, async_upload=async_upload) if decryption_info: decryption_info.url = content_uri @@ -325,7 +329,13 @@ async def _unlocked_transfer_file_to_matrix( thumbnail = thumbnail.location try: db_file.thumbnail = await transfer_thumbnail_to_matrix( - client, intent, thumbnail, video=file, mime_type=mime_type, encrypt=encrypt + client, + intent, + thumbnail, + video=file, + mime_type=mime_type, + encrypt=encrypt, + async_upload=async_upload, ) except FileIdInvalidError: log.warning(f"Failed to transfer thumbnail for {thumbnail!s}", exc_info=True) @@ -340,6 +350,7 @@ async def _unlocked_transfer_file_to_matrix( mime_type=converted_anim.thumbnail_mime, width=converted_anim.width, height=converted_anim.height, + async_upload=async_upload, ) try: diff --git a/mautrix_telegram/util/parallel_file_transfer.py b/mautrix_telegram/util/parallel_file_transfer.py index 42250371..da933c5c 100644 --- a/mautrix_telegram/util/parallel_file_transfer.py +++ b/mautrix_telegram/util/parallel_file_transfer.py @@ -317,6 +317,7 @@ async def parallel_transfer_to_matrix( filename: str, encrypt: bool, parallel_id: int, + async_upload: bool = False, ) -> DBTelegramFile: size = location.size mime_type = location.mime_type @@ -340,7 +341,11 @@ async def parallel_transfer_to_matrix( data = encrypted(data) up_mime_type = "application/octet-stream" content_uri = await intent.upload_media( - data, mime_type=up_mime_type, filename=filename, size=size if not encrypt else None + data, + mime_type=up_mime_type, + filename=filename, + size=size if not encrypt else None, + async_upload=async_upload and not encrypt, ) if decryption_info: decryption_info.url = content_uri diff --git a/requirements.txt b/requirements.txt index edf66f7f..a9799c37 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ python-magic>=0.4,<0.5 commonmark>=0.8,<0.10 aiohttp>=3,<4 yarl>=1,<2 -mautrix>=0.15.1,<0.16 +mautrix>=0.15.4,<0.16 #telethon>=1.24,<1.25 # Fork to make session storage async and update to layer 138 tulir-telethon==1.25.0a5