Skip to content

Commit a0bb0b5

Browse files
null-nicknull-nick
committed
1 parent b6f49d0 commit a0bb0b5

2 files changed

Lines changed: 103 additions & 94 deletions

File tree

‎pyrogram/client.py‎

Lines changed: 102 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -32,15 +32,16 @@
3232
from io import StringIO, BytesIO
3333
from mimetypes import MimeTypes
3434
from pathlib import Path
35-
from typing import Union, List, Optional, Callable, AsyncGenerator, Type
35+
from typing import Union, List, Optional, Callable, AsyncGenerator, Type, Tuple
3636

3737
import pyrogram
3838
from pyrogram import __version__, __license__
3939
from pyrogram import enums
4040
from pyrogram import raw
4141
from pyrogram import utils
4242
from pyrogram.crypto import aes
43-
from pyrogram.errors import CDNFileHashMismatch, AuthBytesInvalid
43+
from pyrogram.errors import CDNFileHashMismatch, AuthBytesInvalid, ChannelInvalid, PersistentTimestampInvalid, \
44+
PersistentTimestampOutdated
4445
from pyrogram.errors import (
4546
SessionPasswordNeeded,
4647
VolumeLocNotFound, ChannelPrivate,
@@ -186,7 +187,7 @@ class Client(Methods):
186187
Defaults to False, because ``getpass`` (the library used) is known to be problematic in some
187188
terminal environments.
188189
189-
max_concurrent_transmissions (``bool``, *optional*):
190+
max_concurrent_transmissions (``int``, *optional*):
190191
Set the maximum amount of concurrent transmissions (uploads & downloads).
191192
A value that is too high may result in network related issues.
192193
Defaults to 1.
@@ -576,14 +577,14 @@ async def handle_updates(self, updates):
576577
pts = getattr(update, "pts", None)
577578
pts_count = getattr(update, "pts_count", None)
578579

579-
if pts:
580+
if pts and not self.skip_updates:
580581
await self.storage.update_state(
581582
(
582-
utils.get_channel_id(channel_id) if channel_id else self.me.id,
583+
utils.get_channel_id(channel_id) if channel_id else 0,
583584
pts,
584585
None,
585586
updates.date,
586-
None
587+
updates.seq
587588
)
588589
)
589590

@@ -617,15 +618,16 @@ async def handle_updates(self, updates):
617618

618619
self.dispatcher.updates_queue.put_nowait((update, users, chats))
619620
elif isinstance(updates, (raw.types.UpdateShortMessage, raw.types.UpdateShortChatMessage)):
620-
await self.storage.update_state(
621-
(
622-
self.me.id,
623-
updates.pts,
624-
None,
625-
updates.date,
626-
None
621+
if not self.skip_updates:
622+
await self.storage.update_state(
623+
(
624+
0,
625+
updates.pts,
626+
None,
627+
updates.date,
628+
None
629+
)
627630
)
628-
)
629631

630632
diff = await self.invoke(
631633
raw.functions.updates.GetDifference(
@@ -653,6 +655,92 @@ async def handle_updates(self, updates):
653655
elif isinstance(updates, raw.types.UpdatesTooLong):
654656
log.info(updates)
655657

658+
async def recover_gaps(self) -> Tuple[int, int]:
659+
states = await self.storage.update_state()
660+
661+
message_updates_counter = 0
662+
other_updates_counter = 0
663+
664+
if not states:
665+
log.info("No states found, skipping recovery.")
666+
return message_updates_counter, other_updates_counter
667+
668+
for state in states:
669+
id, local_pts, _, local_date, _ = state
670+
671+
prev_pts = 0
672+
673+
while True:
674+
try:
675+
diff = await self.invoke(
676+
raw.functions.updates.GetChannelDifference(
677+
channel=await self.resolve_peer(id),
678+
filter=raw.types.ChannelMessagesFilterEmpty(),
679+
pts=local_pts,
680+
limit=10000,
681+
force=False
682+
) if id < 0 else
683+
raw.functions.updates.GetDifference(
684+
pts=local_pts,
685+
date=local_date,
686+
qts=0
687+
)
688+
)
689+
except (ChannelPrivate, ChannelInvalid, PersistentTimestampOutdated, PersistentTimestampInvalid):
690+
break
691+
692+
if isinstance(diff, raw.types.updates.DifferenceEmpty):
693+
break
694+
elif isinstance(diff, raw.types.updates.DifferenceTooLong):
695+
break
696+
elif isinstance(diff, raw.types.updates.Difference):
697+
local_pts = diff.state.pts
698+
elif isinstance(diff, raw.types.updates.DifferenceSlice):
699+
local_pts = diff.intermediate_state.pts
700+
local_date = diff.intermediate_state.date
701+
702+
if prev_pts == local_pts:
703+
break
704+
705+
prev_pts = local_pts
706+
elif isinstance(diff, raw.types.updates.ChannelDifferenceEmpty):
707+
break
708+
elif isinstance(diff, raw.types.updates.ChannelDifferenceTooLong):
709+
break
710+
elif isinstance(diff, raw.types.updates.ChannelDifference):
711+
local_pts = diff.pts
712+
713+
users = {i.id: i for i in diff.users}
714+
chats = {i.id: i for i in diff.chats}
715+
716+
for message in diff.new_messages:
717+
message_updates_counter += 1
718+
self.dispatcher.updates_queue.put_nowait(
719+
(
720+
raw.types.UpdateNewMessage(
721+
message=message,
722+
pts=local_pts,
723+
pts_count=-1
724+
),
725+
users,
726+
chats
727+
)
728+
)
729+
730+
for update in diff.other_updates:
731+
other_updates_counter += 1
732+
self.dispatcher.updates_queue.put_nowait(
733+
(update, users, chats)
734+
)
735+
736+
if isinstance(diff, (raw.types.updates.Difference, raw.types.updates.ChannelDifference)):
737+
break
738+
739+
await self.storage.update_state(id)
740+
741+
log.info("Recovered %s messages and %s updates.", message_updates_counter, other_updates_counter)
742+
return message_updates_counter, other_updates_counter
743+
656744
async def load_session(self):
657745
await self.storage.open()
658746

‎pyrogram/dispatcher.py‎

Lines changed: 1 addition & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323

2424
import pyrogram
2525
from pyrogram import utils
26-
from pyrogram import raw
2726
from pyrogram.handlers import (
2827
CallbackQueryHandler, MessageHandler, EditedMessageHandler, DeletedMessagesHandler,
2928
UserStatusHandler, RawUpdateHandler, InlineQueryHandler, PollHandler,
@@ -155,85 +154,7 @@ async def start(self):
155154
log.info("Started %s HandlerTasks", self.client.workers)
156155

157156
if not self.client.skip_updates:
158-
states = await self.client.storage.update_state()
159-
160-
if not states:
161-
log.info("No states found, skipping recovery.")
162-
return
163-
164-
message_updates_counter = 0
165-
other_updates_counter = 0
166-
167-
for state in states:
168-
id, local_pts, _, local_date, _ = state
169-
170-
prev_pts = 0
171-
172-
while True:
173-
diff = await self.client.invoke(
174-
raw.functions.updates.GetDifference(
175-
pts=local_pts,
176-
date=local_date,
177-
qts=0
178-
) if id == self.client.me.id else
179-
raw.functions.updates.GetChannelDifference(
180-
channel=await self.client.resolve_peer(id),
181-
filter=raw.types.ChannelMessagesFilterEmpty(),
182-
pts=local_pts,
183-
limit=10000
184-
)
185-
)
186-
187-
if isinstance(diff, (raw.types.updates.DifferenceEmpty, raw.types.updates.ChannelDifferenceEmpty)):
188-
break
189-
elif isinstance(diff, (raw.types.updates.DifferenceTooLong, raw.types.updates.ChannelDifferenceTooLong)):
190-
break
191-
elif isinstance(diff, raw.types.updates.ChannelDifference):
192-
local_pts = diff.pts
193-
elif isinstance(diff, raw.types.updates.Difference):
194-
local_pts = diff.state.pts
195-
elif isinstance(diff, raw.types.updates.DifferenceSlice):
196-
local_pts = diff.intermediate_state.pts
197-
local_date = diff.intermediate_state.date
198-
199-
if prev_pts == local_pts:
200-
break
201-
202-
prev_pts = local_pts
203-
204-
users = {i.id: i for i in diff.users}
205-
chats = {i.id: i for i in diff.chats}
206-
207-
for message in diff.new_messages:
208-
message_updates_counter += 1
209-
self.updates_queue.put_nowait(
210-
(
211-
raw.types.UpdateNewMessage(
212-
message=message,
213-
pts=local_pts,
214-
pts_count=-1
215-
) if id == self.client.me.id else
216-
raw.types.UpdateNewChannelMessage(
217-
message=message,
218-
pts=local_pts,
219-
pts_count=-1
220-
),
221-
users,
222-
chats
223-
)
224-
)
225-
226-
for update in diff.other_updates:
227-
other_updates_counter += 1
228-
self.updates_queue.put_nowait(
229-
(update, users, chats)
230-
)
231-
232-
if isinstance(diff, (raw.types.updates.Difference, raw.types.updates.ChannelDifference)):
233-
break
234-
235-
await self.client.storage.update_state(None)
236-
log.info("Recovered %s messages and %s updates.", message_updates_counter, other_updates_counter)
157+
await self.client.recover_gaps()
237158

238159
async def stop(self):
239160
if not self.client.no_updates:

0 commit comments

Comments
 (0)