From 68ea208b4378263d2e04abd8357fe957770ec620 Mon Sep 17 00:00:00 2001
From: Kacnep89 <76146578+Kacnep89@users.noreply.github.com>
Date: Tue, 28 Mar 2023 22:00:36 +0500
Subject: [PATCH] Periodically save update state (#4071)

---
 telethon/client/telegrambaseclient.py | 25 ++++++++++++++-----------
 telethon/client/updates.py            |  2 ++
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py
index 41343e07..54c02a16 100644
--- a/telethon/client/telegrambaseclient.py
+++ b/telethon/client/telegrambaseclient.py
@@ -654,6 +654,19 @@ class TelegramBaseClient(abc.ABC):
             else:
                 connection._proxy = proxy
 
+    def _save_states_and_entities(self: 'TelegramClient'):
+        entities = self._mb_entity_cache.get_all_entities()
+
+        # Piggy-back on an arbitrary TL type with users and chats so the session can understand to read the entities.
+        # It doesn't matter if we put users in the list of chats.
+        self.session.process_entities(types.contacts.ResolvedPeer(None, [e._as_input_peer() for e in entities], []))
+
+        ss, cs = self._message_box.session_state()
+        self.session.set_update_state(0, types.updates.State(**ss, unread_count=0))
+        now = datetime.datetime.now()  # any datetime works; channels don't need it
+        for channel_id, pts in cs.items():
+            self.session.set_update_state(channel_id, types.updates.State(pts, 0, now, 0, unread_count=0))
+
     async def _disconnect_coro(self: 'TelegramClient'):
         if self.session is None:
             return  # already logged out and disconnected
@@ -684,17 +697,7 @@ class TelegramBaseClient(abc.ABC):
             await asyncio.wait(self._event_handler_tasks)
             self._event_handler_tasks.clear()
 
-        entities = self._mb_entity_cache.get_all_entities()
-
-        # Piggy-back on an arbitrary TL type with users and chats so the session can understand to read the entities.
-        # It doesn't matter if we put users in the list of chats.
-        self.session.process_entities(types.contacts.ResolvedPeer(None, [e._as_input_peer() for e in entities], []))
-
-        ss, cs = self._message_box.session_state()
-        self.session.set_update_state(0, types.updates.State(**ss, unread_count=0))
-        now = datetime.datetime.now()  # any datetime works; channels don't need it
-        for channel_id, pts in cs.items():
-            self.session.set_update_state(channel_id, types.updates.State(pts, 0, now, 0, unread_count=0))
+        self._save_states_and_entities()
 
         self.session.close()
 
diff --git a/telethon/client/updates.py b/telethon/client/updates.py
index 157d1f47..8cc3a04b 100644
--- a/telethon/client/updates.py
+++ b/telethon/client/updates.py
@@ -467,6 +467,8 @@ class UpdateMethods:
             # inserted because this is a rather expensive operation
             # (default's sqlite3 takes ~0.1s to commit changes). Do
             # it every minute instead. No-op if there's nothing new.
+            self._save_states_and_entities()
+
             self.session.save()
 
     async def _dispatch_update(self: 'TelegramClient', update):