From 2f4065c958e7bf5f115e2b3bc9cf215289bfa1c1 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Fri, 1 Mar 2024 21:49:44 +0100 Subject: [PATCH] Improve session storage lifecycle --- client/src/telethon/_impl/client/client/auth.py | 7 ++++++- client/src/telethon/_impl/client/client/net.py | 9 +++++++-- .../src/telethon/_impl/session/chat/hash_cache.py | 5 +++++ .../_impl/session/message_box/messagebox.py | 8 ++++++++ .../src/telethon/_impl/session/storage/memory.py | 4 ++-- .../src/telethon/_impl/session/storage/sqlite.py | 5 +---- .../src/telethon/_impl/session/storage/storage.py | 15 +++++++++------ 7 files changed, 38 insertions(+), 15 deletions(-) diff --git a/client/src/telethon/_impl/client/client/auth.py b/client/src/telethon/_impl/client/client/auth.py index 14666734..e9fc1da2 100644 --- a/client/src/telethon/_impl/client/client/auth.py +++ b/client/src/telethon/_impl/client/client/auth.py @@ -260,4 +260,9 @@ async def check_password( async def sign_out(self: Client) -> None: await self(functions.auth.log_out()) - await self._storage.delete() + self._chat_hashes.clear() + self._message_box.reset() + + self._session.user = None + self._session.state = None + await self._storage.save(self._session) diff --git a/client/src/telethon/_impl/client/client/net.py b/client/src/telethon/_impl/client/client/net.py index 2ac8b9c0..3b1b24c2 100644 --- a/client/src/telethon/_impl/client/client/net.py +++ b/client/src/telethon/_impl/client/client/net.py @@ -230,8 +230,13 @@ async def disconnect(self: Client) -> None: "unhandled exception during disconnect; this is a bug" ) - self._session.state = self._message_box.session_state() - await self._storage.save(self._session) + try: + if self._session.user: + # Only save if we haven't logged out (prevents double-save) + self._session.state = self._message_box.session_state() + await self._storage.save(self._session) + finally: + await self._storage.close() async def invoke_request( diff --git a/client/src/telethon/_impl/session/chat/hash_cache.py b/client/src/telethon/_impl/session/chat/hash_cache.py index f78feae1..bbdc12b3 100644 --- a/client/src/telethon/_impl/session/chat/hash_cache.py +++ b/client/src/telethon/_impl/session/chat/hash_cache.py @@ -33,6 +33,11 @@ class ChatHashCache: else: return None + def clear(self) -> None: + self._hash_map.clear() + self._self_id = None + self._self_bot = False + def _has(self, id: int) -> bool: return id in self._hash_map diff --git a/client/src/telethon/_impl/session/message_box/messagebox.py b/client/src/telethon/_impl/session/message_box/messagebox.py index 7a35147e..41150601 100644 --- a/client/src/telethon/_impl/session/message_box/messagebox.py +++ b/client/src/telethon/_impl/session/message_box/messagebox.py @@ -103,6 +103,14 @@ class MessageBox: self.getting_diff_for.clear() self.next_deadline = ENTRY_ACCOUNT + def reset(self) -> None: + self.map.clear() + self.date = epoch() + self.seq = NO_SEQ + self.possible_gaps.clear() + self.getting_diff_for.clear() + self.next_deadline = None + def session_state(self) -> UpdateState: return UpdateState( pts=self.map[ENTRY_ACCOUNT].pts if ENTRY_ACCOUNT in self.map else NO_PTS, diff --git a/client/src/telethon/_impl/session/storage/memory.py b/client/src/telethon/_impl/session/storage/memory.py index d7458d9b..a5bcc036 100644 --- a/client/src/telethon/_impl/session/storage/memory.py +++ b/client/src/telethon/_impl/session/storage/memory.py @@ -24,5 +24,5 @@ class MemorySession(Storage): async def save(self, session: Session) -> None: self.session = session - async def delete(self) -> None: - self.session = None + async def close(self) -> None: + pass diff --git a/client/src/telethon/_impl/session/storage/sqlite.py b/client/src/telethon/_impl/session/storage/sqlite.py index 70f682cc..a9d0ce53 100644 --- a/client/src/telethon/_impl/session/storage/sqlite.py +++ b/client/src/telethon/_impl/session/storage/sqlite.py @@ -52,14 +52,11 @@ class SqliteSession(Storage): conn = self._current_conn() with conn: self._save_v10(conn.cursor(), session) - conn.close() - self._conn = None - async def delete(self) -> None: + async def close(self) -> None: if self._conn: self._conn.close() self._conn = None - self._path.unlink() def _current_conn(self) -> sqlite3.Connection: if self._conn is None: diff --git a/client/src/telethon/_impl/session/storage/storage.py b/client/src/telethon/_impl/session/storage/storage.py index fcf56f47..8c7fdc3b 100644 --- a/client/src/telethon/_impl/session/storage/storage.py +++ b/client/src/telethon/_impl/session/storage/storage.py @@ -14,7 +14,7 @@ class Storage(abc.ABC): """ Load the :class:`Session` instance, if any. - This method is called by the library prior to `connect`. + This method is called by the library prior to ``connect``. :return: The previously-saved session. """ @@ -24,18 +24,21 @@ class Storage(abc.ABC): """ Save the :class:`Session` instance to persistent storage. - This method is called by the library post `disconnect`. + This method is called by the library after significant changes to the session, + such as login, logout, or to persist the update state prior to disconnection. :param session: The session information that should be persisted. """ @abc.abstractmethod - async def delete(self) -> None: + async def close(self) -> None: """ - Delete the saved `Session`. + Close the :class:`Session` instance, if it was still open. - This method is called by the library post `log_out`. + This method is called by the library post ``disconnect``, + even if the call to :meth:`save` failed. - Note that both :meth:`load` and :meth:`save` may still be called after. + Note that :meth:`load` may still be called after, + in which case the session should be reopened. """