Persist next_batch for custom puppets

alembic/versions/a7c04a56041b_store_custom_puppet_next_batch_in_.py

@@ -0,0 +1,26 @@
+"""Store custom puppet next_batch in database
+
+Revision ID: a7c04a56041b
+Revises: 4f7d7ed5792a
+Create Date: 2019-08-06 23:08:51.087651
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "a7c04a56041b"
+down_revision = "4f7d7ed5792a"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.batch_alter_table("puppet") as batch_op:
+        batch_op.add_column(sa.Column("next_batch", sa.String(), nullable=True))
+
+
+def downgrade():
+    with op.batch_alter_table("puppet") as batch_op:
+        batch_op.drop_column("next_batch")

mautrix_telegram/db/puppet.py

@@ -20,7 +20,7 @@ from sqlalchemy.sql import expression
 from sqlalchemy.engine.result import RowProxy
 from sqlalchemy.sql.expression import ClauseElement
 
-from mautrix.types import UserID
+from mautrix.types import UserID, SyncToken
 from mautrix.bridge.db import Base
 
 from ..types import TelegramID
@@ -32,6 +32,7 @@ class Puppet(Base):
     id: TelegramID = Column(Integer, primary_key=True)
     custom_mxid: UserID = Column(String, nullable=True)
     access_token: str = Column(String, nullable=True)
+    next_batch: SyncToken = Column(String, nullable=True)
     displayname: str = Column(String, nullable=True)
     displayname_source: TelegramID = Column(Integer, nullable=True)
     username: str = Column(String, nullable=True)
@@ -42,12 +43,12 @@ class Puppet(Base):
 
     @classmethod
     def scan(cls, row: RowProxy) -> Optional['Puppet']:
-        (id, custom_mxid, access_token, displayname, displayname_source, username, photo_id,
+        (id, custom_mxid, access_token, next_batch, displayname, displayname_source, username,
-         is_bot, matrix_registered, disable_updates) = row
+         photo_id, is_bot, matrix_registered, disable_updates) = row
-        return cls(id=id, custom_mxid=custom_mxid, access_token=access_token,
+        return cls(id=id, custom_mxid=custom_mxid, access_token=access_token, username=username,
-                   displayname=displayname, displayname_source=displayname_source,
+                   next_batch=next_batch, displayname=displayname, photo_id=photo_id,
-                   username=username, photo_id=photo_id, is_bot=is_bot,
+                   displayname_source=displayname_source, matrix_registered=matrix_registered,
-                   matrix_registered=matrix_registered, disable_updates=disable_updates)
+                   disable_updates=disable_updates, is_bot=is_bot)
 
     @classmethod
     def all_with_custom_mxid(cls) -> Iterable['Puppet']:
@@ -79,6 +80,7 @@ class Puppet(Base):
         with self.db.begin() as conn:
             conn.execute(self.t.insert().values(
                 id=self.id, custom_mxid=self.custom_mxid, access_token=self.access_token,
-                displayname=self.displayname, displayname_source=self.displayname_source,
+                next_batch=self.next_batch, displayname=self.displayname, username=self.username,
-                username=self.username, photo_id=self.photo_id, is_bot=self.is_bot,
+                displayname_source=self.displayname_source, photo_id=self.photo_id,
-                matrix_registered=self.matrix_registered, disable_updates=self.disable_updates))
+                is_bot=self.is_bot, matrix_registered=self.matrix_registered,
+                disable_updates=self.disable_updates))

mautrix_telegram/puppet.py

@@ -24,7 +24,7 @@ from telethon.tl.types import (UserProfilePhoto, User, UpdateUserName, PeerUser,
 from mautrix.appservice import AppService, IntentAPI
 from mautrix.errors import MatrixRequestError
 from mautrix.bridge import CustomPuppetMixin
-from mautrix.types import UserID
+from mautrix.types import UserID, SyncToken
 
 from .types import TelegramID
 from .db import Puppet as DBPuppet
@@ -57,6 +57,7 @@ class Puppet(CustomPuppetMixin):
     id: TelegramID
     access_token: Optional[str]
     custom_mxid: Optional[UserID]
+    _next_batch: Optional[SyncToken]
     default_mxid: UserID
 
     username: Optional[str]
@@ -78,6 +79,7 @@ class Puppet(CustomPuppetMixin):
                  id: TelegramID,
                  access_token: Optional[str] = None,
                  custom_mxid: Optional[UserID] = None,
+                 next_batch: Optional[SyncToken] = None,
                  username: Optional[str] = None,
                  displayname: Optional[str] = None,
                  displayname_source: Optional[TelegramID] = None,
@@ -89,6 +91,7 @@ class Puppet(CustomPuppetMixin):
         self.id = id
         self.access_token = access_token
         self.custom_mxid = custom_mxid
+        self._next_batch = next_batch
         self.default_mxid = self.get_mxid_from_id(self.id)
 
         self.username = username
@@ -118,6 +121,15 @@ class Puppet(CustomPuppetMixin):
     def peer(self) -> PeerUser:
         return PeerUser(user_id=self.tgid)
 
+    @property
+    def next_batch(self) -> SyncToken:
+        return self._next_batch
+
+    @next_batch.setter
+    def next_batch(self, value: SyncToken) -> None:
+        self._next_batch = value
+        self.db_instance.edit(next_batch=self._next_batch)
+
     @staticmethod
     async def is_logged_in() -> bool:
         """ Is True if the puppet is logged in. """
@@ -148,26 +160,27 @@ class Puppet(CustomPuppetMixin):
             self._db_instance = self.new_db_instance()
         return self._db_instance
 
+    @property
+    def _fields(self) -> Dict[str, Any]:
+        return dict(access_token=self.access_token, next_batch=self._next_batch,
+                    custom_mxid=self.custom_mxid, username=self.username, is_bot=self.is_bot,
+                    displayname=self.displayname, displayname_source=self.displayname_source,
+                    photo_id=self.photo_id, matrix_registered=self.is_registered,
+                    disable_updates=self.disable_updates)
+
     def new_db_instance(self) -> DBPuppet:
-        return DBPuppet(id=self.id, access_token=self.access_token, custom_mxid=self.custom_mxid,
+        return DBPuppet(id=self.id, **self._fields)
-                        username=self.username, displayname=self.displayname,
+
-                        displayname_source=self.displayname_source, photo_id=self.photo_id,
+    def save(self) -> None:
-                        is_bot=self.is_bot, matrix_registered=self.is_registered,
+        self.db_instance.edit(**self._fields)
-                        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.next_batch, db_puppet.username, db_puppet.displayname,
-                      db_puppet.photo_id, db_puppet.is_bot, db_puppet.matrix_registered,
+                      db_puppet.displayname_source, db_puppet.photo_id, db_puppet.is_bot,
-                      db_puppet.disable_updates, db_instance=db_puppet)
+                      db_puppet.matrix_registered, db_puppet.disable_updates,
-
+                      db_instance=db_puppet)
-    def save(self) -> None:
-        self.db_instance.edit(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,
-                              disable_updates=self.disable_updates)
 
     # endregion
     # region Info updating

mautrix_telegram/user.py

@@ -13,15 +13,14 @@
 #
 # 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 (Awaitable, Dict, List, Iterable, Match, NewType, Optional, Tuple, Any, Union,
+from typing import (Awaitable, Dict, List, Iterable, Match, NewType, Optional, Tuple, Any,
-                    AsyncIterable, TYPE_CHECKING)
+                    TYPE_CHECKING)
 import logging
 import asyncio
 import re
 
 from telethon.tl.types import (TypeUpdate, UpdateNewMessage, UpdateNewChannelMessage, PeerUser,
-                               UpdateShortChatMessage, UpdateShortMessage, User as TLUser,
+                               UpdateShortChatMessage, UpdateShortMessage, User as TLUser, Chat)
-                               ChannelForbidden, ChatForbidden, Chat, Channel)
 from telethon.tl.types.contacts import ContactsNotModified
 from telethon.tl.functions.contacts import GetContactsRequest, SearchRequest
 from telethon.tl.functions.account import UpdateStatusRequest