Persist updates.State upon disconnection

telethon/sessions/abstract.py

@@ -67,6 +67,25 @@ class Session(ABC):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def get_update_state(self, entity_id):
+        """
+        Returns the ``UpdateState`` associated with the given `entity_id`.
+        If the `entity_id` is 0, it should return the ``UpdateState`` for
+        no specific channel (the "general" state). If no state is known
+        it should ``return None``.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def set_update_state(self, entity_id, state):
+        """
+        Sets the given ``UpdateState`` for the specified `entity_id`, which
+        should be 0 if the ``UpdateState`` is the "general" state (and not
+        for any specific channel).
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def close(self):
         """

telethon/sessions/memory.py

@@ -35,6 +35,7 @@ class MemorySession(Session):
 
         self._files = {}
         self._entities = set()
+        self._update_states = {}
 
     def set_dc(self, dc_id, server_address, port):
         self._dc_id = dc_id or 0
@@ -57,6 +58,12 @@ class MemorySession(Session):
     def auth_key(self, value):
         self._auth_key = value
 
+    def get_update_state(self, entity_id):
+        return self._update_states.get(entity_id, None)
+
+    def set_update_state(self, entity_id, state):
+        self._update_states[entity_id] = state
+
     def close(self):
         pass
 

telethon/sessions/sqlite.py

@@ -5,6 +5,8 @@ from base64 import b64decode
 from os.path import isfile as file_exists
 from threading import Lock, RLock
 
+from telethon.tl import types
+
 from .memory import MemorySession, _SentFileType
 from .. import utils
 from ..crypto import AuthKey
@@ -226,6 +228,22 @@ class SQLiteSession(MemorySession):
             ))
             c.close()
 
+    def get_update_state(self, entity_id):
+        c = self._cursor()
+        row = c.execute('select pts, qts, date, seq from update_state '
+                        'where id = ?', (entity_id,)).fetchone()
+        c.close()
+        if row:
+            return types.updates.State(*row)
+
+    def set_update_state(self, entity_id, state):
+        with self._db_lock:
+            c = self._cursor()
+            c.execute('insert or replace into update_state values (?,?,?,?,?)',
+                      (entity_id, state.pts, state.qts, state.date, state.seq))
+            c.close()
+            self.save()
+
     def save(self):
         """Saves the current session object as session_user_id.session"""
         with self._db_lock: