From 2a00bcaa1222fe746b4a983388a711269aa38ba9 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Wed, 25 Apr 2018 13:37:29 +0200
Subject: [PATCH] Persist updates.State upon disconnection

---
 telethon/sessions/abstract.py    | 19 +++++++++++++++++++
 telethon/sessions/memory.py      |  7 +++++++
 telethon/sessions/sqlite.py      | 18 ++++++++++++++++++
 telethon/telegram_bare_client.py |  1 +
 telethon/update_state.py         |  4 ++++
 5 files changed, 49 insertions(+)

diff --git a/telethon/sessions/abstract.py b/telethon/sessions/abstract.py
index 75324077..fb0778bd 100644
--- a/telethon/sessions/abstract.py
+++ b/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):
         """
diff --git a/telethon/sessions/memory.py b/telethon/sessions/memory.py
index 330a82c1..acd09a77 100644
--- a/telethon/sessions/memory.py
+++ b/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
 
diff --git a/telethon/sessions/sqlite.py b/telethon/sessions/sqlite.py
index 1b7d3c31..7e7380af 100644
--- a/telethon/sessions/sqlite.py
+++ b/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:
diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py
index d1f4333f..826cc155 100644
--- a/telethon/telegram_bare_client.py
+++ b/telethon/telegram_bare_client.py
@@ -275,6 +275,7 @@ class TelegramBareClient:
 
         # TODO Shall we clear the _exported_sessions, or may be reused?
         self._first_request = True  # On reconnect it will be first again
+        self.session.set_update_state(0, self.updates.get_update_state(0))
         self.session.close()
 
     def _reconnect(self, new_dc=None):
diff --git a/telethon/update_state.py b/telethon/update_state.py
index e679e99e..eaedcfc3 100644
--- a/telethon/update_state.py
+++ b/telethon/update_state.py
@@ -110,6 +110,10 @@ class UpdateState:
                 # We don't want to crash a worker thread due to any reason
                 __log__.exception('Unhandled exception on worker %d', wid)
 
+    def get_update_state(self, entity_id):
+        """Gets the updates.State corresponding to the given entity or 0."""
+        return self._state
+
     def process(self, update):
         """Processes an update object. This method is normally called by
            the library itself.