From e6ef40e51dc9ef110ded0b60a17b4433e3d9c1be Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Thu, 19 May 2022 13:15:39 +0300
Subject: [PATCH] Update Telethon

---
 mautrix_telegram/abstract_user.py       |  3 +++
 mautrix_telegram/db/telethon_session.py | 18 +++++++++---------
 requirements.txt                        |  3 +--
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/mautrix_telegram/abstract_user.py b/mautrix_telegram/abstract_user.py
index f367a671..b4be3f90 100644
--- a/mautrix_telegram/abstract_user.py
+++ b/mautrix_telegram/abstract_user.py
@@ -57,6 +57,7 @@ from telethon.tl.types import (
     UpdateReadChannelInbox,
     UpdateReadHistoryInbox,
     UpdateReadHistoryOutbox,
+    UpdateShort,
     UpdateShortChatMessage,
     UpdateShortMessage,
     UpdateUserName,
@@ -297,6 +298,8 @@ class AbstractUser(ABC):
     # region Telegram update handling
 
     async def _update(self, update: TypeUpdate) -> None:
+        if isinstance(update, UpdateShort):
+            update = update.update
         asyncio.create_task(self._handle_entity_updates(getattr(update, "_entities", {})))
         if isinstance(
             update,
diff --git a/mautrix_telegram/db/telethon_session.py b/mautrix_telegram/db/telethon_session.py
index a3d139f9..67d3b315 100644
--- a/mautrix_telegram/db/telethon_session.py
+++ b/mautrix_telegram/db/telethon_session.py
@@ -15,7 +15,7 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, ClassVar, Iterable
 import asyncio
 import datetime
 
@@ -124,19 +124,19 @@ class PgSession(MemorySession):
         return updates.State(row["pts"], row["qts"], date, row["seq"], row["unread_count"])
 
     async def set_update_state(self, entity_id: int, row: updates.State) -> None:
-        q = (
-            "INSERT INTO telethon_update_state"
-            "    (session_id, entity_id, pts, qts, date, seq, unread_count) "
-            "VALUES ($1, $2, $3, $4, $5, $6, $7)"
-            "ON CONFLICT (session_id, entity_id) DO UPDATE"
-            "    SET pts=$3, qts=$4, date=$5, seq=$6, unread_count=$7"
-        )
+        q = """
+        INSERT INTO telethon_update_state(session_id, entity_id, pts, qts, date, seq, unread_count)
+        VALUES ($1, $2, $3, $4, $5, $6, $7)
+        ON CONFLICT (session_id, entity_id) DO UPDATE SET
+            pts=excluded.pts, qts=excluded.qts, date=excluded.date, seq=excluded.seq,
+            unread_count=excluded.unread_count
+        """
         ts = row.date.timestamp()
         await self.db.execute(
             q, self.session_id, entity_id, row.pts, row.qts, ts, row.seq, row.unread_count
         )
 
-    async def get_update_states(self) -> tuple[tuple[int, updates.State], ...]:
+    async def get_update_states(self) -> Iterable[tuple[int, updates.State], ...]:
         q = (
             "SELECT entity_id, pts, qts, date, seq, unread_count FROM telethon_update_state "
             "WHERE session_id=$1"
diff --git a/requirements.txt b/requirements.txt
index 3f86111b..4af94dbc 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,8 +5,7 @@ aiohttp>=3,<4
 yarl>=1,<2
 mautrix>=0.16.4,<0.17
 #telethon>=1.24,<1.25
-# Fork to make session storage async and update to layer 138
-tulir-telethon==1.25.0a7
+tulir-telethon==1.25.0a8
 asyncpg>=0.20,<0.26
 mako>=1,<2
 setuptools