diff --git a/alembic/versions/17574c57f3f8_add_disable_updates_field_for_puppets.py b/alembic/versions/17574c57f3f8_add_disable_updates_field_for_puppets.py new file mode 100644 index 00000000..b91b0f13 --- /dev/null +++ b/alembic/versions/17574c57f3f8_add_disable_updates_field_for_puppets.py @@ -0,0 +1,26 @@ +"""Add disable_updates field for puppets + +Revision ID: 17574c57f3f8 +Revises: a9119be92164 +Create Date: 2019-05-15 00:24:46.967529 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '17574c57f3f8' +down_revision = 'a9119be92164' +branch_labels = None +depends_on = None + + +def upgrade(): + op.add_column("puppet", sa.Column("disable_updates", sa.Boolean(), nullable=False, + server_default=sa.sql.expression.false())) + + +def downgrade(): + with op.batch_alter_table("puppet") as batch_op: + batch_op.drop_column("disable_updates") diff --git a/mautrix_telegram/abstract_user.py b/mautrix_telegram/abstract_user.py index a4c94a61..c188874f 100644 --- a/mautrix_telegram/abstract_user.py +++ b/mautrix_telegram/abstract_user.py @@ -14,7 +14,7 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -from typing import Tuple, Optional, List, Union, TYPE_CHECKING +from typing import Tuple, Optional, List, Union, Dict, TYPE_CHECKING from abc import ABC, abstractmethod import asyncio import logging @@ -201,6 +201,8 @@ class AbstractUser(ABC): # region Telegram update handling async def _update(self, update: TypeUpdate) -> None: + asyncio.ensure_future(self._handle_entity_updates(getattr(update, "_entities", {})), + loop=self.loop) if isinstance(update, (UpdateShortChatMessage, UpdateShortMessage, UpdateNewChannelMessage, UpdateNewMessage, UpdateEditMessage, UpdateEditChannelMessage)): await self.update_message(update) @@ -277,6 +279,15 @@ class AbstractUser(ABC): sender = pu.Puppet.get(TelegramID(update.user_id)) await portal.handle_telegram_typing(sender, update) + async def _handle_entity_updates(self, entities: Dict[int, Union[User, Chat, Channel]]) -> None: + try: + users = (entity for entity in entities.items() if isinstance(entity, User)) + puppets = ((pu.Puppet.get(TelegramID(user.id)), user) for user in users) + await asyncio.gather(*[puppet.update_info(self, info) + for puppet, info in puppets if puppet]) + except Exception: + self.log.exception("Failed to handle entity updates") + async def update_others_info(self, update: Union[UpdateUserName, UpdateUserPhoto]) -> None: # TODO duplication not checked puppet = pu.Puppet.get(TelegramID(update.user_id)) diff --git a/mautrix_telegram/db/puppet.py b/mautrix_telegram/db/puppet.py index 531e7122..489ef672 100644 --- a/mautrix_telegram/db/puppet.py +++ b/mautrix_telegram/db/puppet.py @@ -35,15 +35,16 @@ class Puppet(Base): photo_id = Column(String, nullable=True) is_bot = Column(Boolean, nullable=True) matrix_registered = Column(Boolean, nullable=False, server_default=expression.false()) + disable_updates = Column(Boolean, nullable=False, server_default=expression.false()) @classmethod def scan(cls, row) -> Optional['Puppet']: (id, custom_mxid, access_token, displayname, displayname_source, username, photo_id, - is_bot, matrix_registered) = row + is_bot, matrix_registered, disable_updates) = row return cls(id=id, custom_mxid=custom_mxid, access_token=access_token, displayname=displayname, displayname_source=displayname_source, username=username, photo_id=photo_id, is_bot=is_bot, - matrix_registered=matrix_registered) + matrix_registered=matrix_registered, disable_updates=disable_updates) @classmethod def _one_or_none(cls, rows: RowProxy) -> Optional['Puppet']: @@ -84,4 +85,4 @@ class Puppet(Base): id=self.id, custom_mxid=self.custom_mxid, access_token=self.access_token, displayname=self.displayname, displayname_source=self.displayname_source, username=self.username, photo_id=self.photo_id, is_bot=self.is_bot, - matrix_registered=self.matrix_registered)) + matrix_registered=self.matrix_registered, disable_updates=self.disable_updates)) diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index e1ec9ebc..cb7d144b 100644 --- a/mautrix_telegram/puppet.py +++ b/mautrix_telegram/puppet.py @@ -61,6 +61,7 @@ class Puppet: photo_id: Optional[str] = None, is_bot: bool = False, is_registered: bool = False, + disable_updates: bool = False, db_instance: Optional[DBPuppet] = None) -> None: self.id = id # type: TelegramID self.access_token = access_token # type: Optional[str] @@ -73,6 +74,7 @@ class Puppet: self.photo_id = photo_id # type: Optional[str] self.is_bot = is_bot # type: bool self.is_registered = is_registered # type: bool + self.disable_updates = disable_updates # type: bool self._db_instance = db_instance # type: Optional[DBPuppet] self.default_mxid_intent = self.az.intent.user(self.default_mxid) @@ -281,23 +283,26 @@ class Puppet: return DBPuppet(id=self.id, access_token=self.access_token, custom_mxid=self.custom_mxid, username=self.username, displayname=self.displayname, displayname_source=self.displayname_source, photo_id=self.photo_id, - is_bot=self.is_bot, matrix_registered=self.is_registered) + is_bot=self.is_bot, matrix_registered=self.is_registered, + disable_updates=self.disable_updates) @classmethod def from_db(cls, db_puppet: DBPuppet) -> 'Puppet': return Puppet(db_puppet.id, db_puppet.access_token, db_puppet.custom_mxid, db_puppet.username, db_puppet.displayname, db_puppet.displayname_source, db_puppet.photo_id, db_puppet.is_bot, db_puppet.matrix_registered, - db_instance=db_puppet) + db_puppet.disable_updates, db_instance=db_puppet) def save(self) -> None: self.db_instance.update(access_token=self.access_token, custom_mxid=self.custom_mxid, username=self.username, displayname=self.displayname, displayname_source=self.displayname_source, photo_id=self.photo_id, - is_bot=self.is_bot, matrix_registered=self.is_registered) + is_bot=self.is_bot, matrix_registered=self.is_registered, + disable_updates=self.disable_updates) # endregion # region Info updating + def similarity(self, query: str) -> int: username_similarity = (SequenceMatcher(None, self.username, query).ratio() if self.username else 0) @@ -334,6 +339,8 @@ class Puppet: displayname=name) async def update_info(self, source: 'AbstractUser', info: User) -> None: + if self.disable_updates: + return changed = False if self.username != info.username: self.username = info.username @@ -350,6 +357,8 @@ class Puppet: async def update_displayname(self, source: 'AbstractUser', info: Union[User, UpdateUserName] ) -> bool: + if self.disable_updates: + return False ignore_source = (not source.is_relaybot and self.displayname_source is not None and self.displayname_source != source.tgid) @@ -370,6 +379,8 @@ class Puppet: return False async def update_avatar(self, source: 'AbstractUser', photo: FileLocation) -> bool: + if self.disable_updates: + return False photo_id = f"{photo.volume_id}-{photo.local_id}" if self.photo_id != photo_id: file = await util.transfer_file_to_matrix(source.client, self.default_mxid_intent,