diff --git a/alembic/versions/a7c04a56041b_store_custom_puppet_next_batch_in_.py b/alembic/versions/a7c04a56041b_store_custom_puppet_next_batch_in_.py new file mode 100644 index 00000000..4ea75c6e --- /dev/null +++ b/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") diff --git a/mautrix_telegram/db/puppet.py b/mautrix_telegram/db/puppet.py index bc22d09c..8b3027e2 100644 --- a/mautrix_telegram/db/puppet.py +++ b/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, - 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, disable_updates=disable_updates) + (id, custom_mxid, access_token, next_batch, displayname, displayname_source, username, + photo_id, is_bot, matrix_registered, disable_updates) = row + return cls(id=id, custom_mxid=custom_mxid, access_token=access_token, username=username, + next_batch=next_batch, displayname=displayname, photo_id=photo_id, + displayname_source=displayname_source, matrix_registered=matrix_registered, + 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, - username=self.username, photo_id=self.photo_id, is_bot=self.is_bot, - matrix_registered=self.matrix_registered, disable_updates=self.disable_updates)) + next_batch=self.next_batch, displayname=self.displayname, username=self.username, + displayname_source=self.displayname_source, photo_id=self.photo_id, + is_bot=self.is_bot, matrix_registered=self.matrix_registered, + disable_updates=self.disable_updates)) diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index 7f9e12a5..013c2e42 100644 --- a/mautrix_telegram/puppet.py +++ b/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, - 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) + return DBPuppet(id=self.id, **self._fields) + + def save(self) -> None: + self.db_instance.edit(**self._fields) @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_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) + db_puppet.next_batch, db_puppet.username, db_puppet.displayname, + db_puppet.displayname_source, db_puppet.photo_id, db_puppet.is_bot, + db_puppet.matrix_registered, db_puppet.disable_updates, + db_instance=db_puppet) # endregion # region Info updating diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py index aece4514..324dc3e1 100644 --- a/mautrix_telegram/user.py +++ b/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 . -from typing import (Awaitable, Dict, List, Iterable, Match, NewType, Optional, Tuple, Any, Union, - AsyncIterable, TYPE_CHECKING) +from typing import (Awaitable, Dict, List, Iterable, Match, NewType, Optional, Tuple, Any, + TYPE_CHECKING) import logging import asyncio import re from telethon.tl.types import (TypeUpdate, UpdateNewMessage, UpdateNewChannelMessage, PeerUser, - UpdateShortChatMessage, UpdateShortMessage, User as TLUser, - ChannelForbidden, ChatForbidden, Chat, Channel) + UpdateShortChatMessage, UpdateShortMessage, User as TLUser, Chat) from telethon.tl.types.contacts import ContactsNotModified from telethon.tl.functions.contacts import GetContactsRequest, SearchRequest from telethon.tl.functions.account import UpdateStatusRequest