Update user info from entities attached to updates

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")

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 <https://www.gnu.org/licenses/>.
-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))

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))

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,