diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py index 11ef5483..9708fe21 100644 --- a/telethon/telegram_bare_client.py +++ b/telethon/telegram_bare_client.py @@ -5,22 +5,25 @@ from os import path # Import some externalized utilities to work with the Telegram types and more from . import helpers as utils -from .errors import RPCError, FloodWaitError +from .errors import RPCError, FloodWaitError, FileMigrateError from .network import authenticator, MtProtoSender, TcpTransport from .utils import get_appropriated_part_size # For sending and receiving requests -from .tl import MTProtoRequest +from .tl import MTProtoRequest, JsonSession from .tl.all_tlobjects import layer from .tl.functions import (InitConnectionRequest, InvokeWithLayerRequest) # Initial request from .tl.functions.help import GetConfigRequest -from .tl.functions.auth import ImportAuthorizationRequest +from .tl.functions.auth import ( + ImportAuthorizationRequest, ExportAuthorizationRequest +) # Easier access for working with media from .tl.functions.upload import ( - GetFileRequest, SaveBigFilePartRequest, SaveFilePartRequest) + GetFileRequest, SaveBigFilePartRequest, SaveFilePartRequest +) # All the types we need to work with from .tl.types import InputFile, InputFileBig @@ -64,6 +67,11 @@ class TelegramBareClient: self._timeout = timeout self._logger = logging.getLogger(__name__) + # Cache "exported" senders 'dc_id: TelegramBareClient' and + # their corresponding sessions not to recreate them all + # the time since it's a (somewhat expensive) process. + self._cached_clients = {} + # These will be set later self.dc_options = None self._sender = None @@ -125,8 +133,6 @@ class TelegramBareClient: )) if exported_auth is not None: - # TODO Don't actually need this for exported authorizations, - # they're only valid on such data center. result = self(GetConfigRequest()) # We're only interested in the DC options, @@ -201,6 +207,54 @@ class TelegramBareClient: return next(dc for dc in self.dc_options if dc.id == dc_id) + def _get_exported_client(self, dc_id, + init_connection=False, + bypass_cache=False): + """Gets a cached exported TelegramBareClient for the desired DC. + + If it's the first time retrieving the TelegramBareClient, the + current authorization is exported to the new DC so that + it can be used there, and the connection is initialized. + + If after using the sender a ConnectionResetError is raised, + this method should be called again with init_connection=True + in order to perform the reconnection. + + If bypass_cache is True, a new client will be exported and + it will not be cached. + """ + # Thanks badoualy/kotlogram on /telegram/api/DefaultTelegramClient.kt + # for clearly showing how to export the authorization! ^^ + client = self._cached_clients.get(dc_id) + if client and not bypass_cache: + if init_connection: + client.reconnect() + return client + else: + dc = self._get_dc(dc_id) + + # Export the current authorization to the new DC. + export_auth = self(ExportAuthorizationRequest(dc_id)) + + # Create a temporary session for this IP address, which needs + # to be different because each auth_key is unique per DC. + # + # Construct this session with the connection parameters + # (system version, device model...) from the current one. + session = JsonSession(self.session) + session.server_address = dc.ip_address + session.port = dc.port + client = TelegramBareClient( + session, self.api_id, self.api_hash, + timeout=self._timeout + ) + client.connect(exported_auth=export_auth) + + if not bypass_cache: + # Don't go through this expensive process every time. + self._cached_clients[dc_id] = client + return client + # endregion # region Invoking Telegram requests @@ -341,12 +395,21 @@ class TelegramBareClient: else: f = file + # The used client will change if FileMigrateError occurs + client = self + try: offset_index = 0 while True: offset = offset_index * part_size - result = self( - GetFileRequest(input_location, offset, part_size)) + + try: + result = client( + GetFileRequest(input_location, offset, part_size)) + except FileMigrateError as e: + client = self._get_exported_client(e.new_dc) + continue + offset_index += 1 # If we have received no data (0 bytes), the file is over diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py index 5d0362b0..64e269aa 100644 --- a/telethon/telegram_client.py +++ b/telethon/telegram_client.py @@ -8,8 +8,8 @@ from . import TelegramBareClient # Import some externalized utilities to work with the Telegram types and more from . import helpers as utils from .errors import (RPCError, UnauthorizedError, InvalidParameterError, - ReadCancelledError, FileMigrateError, PhoneMigrateError, - NetworkMigrateError, UserMigrateError, PhoneCodeEmptyError, + ReadCancelledError, PhoneCodeEmptyError, + PhoneMigrateError, NetworkMigrateError, UserMigrateError, PhoneCodeExpiredError, PhoneCodeHashEmptyError, PhoneCodeInvalidError, InvalidChecksumError) @@ -123,10 +123,6 @@ class TelegramClient(TelegramBareClient): self.session.system_lang_code = \ system_lang_code if system_lang_code else self.session.lang_code - # Cache "exported" senders 'dc_id: MtProtoSender' and - # their corresponding sessions not to recreate them all - # the time since it's a (somewhat expensive) process. - self._cached_clients = {} self._updates_thread = None self._phone_code_hashes = {} @@ -162,55 +158,6 @@ class TelegramClient(TelegramBareClient): # region Working with different connections - def _get_exported_client(self, dc_id, - init_connection=False, - bypass_cache=False): - """Gets a cached exported TelegramBareClient for the desired DC. - - If it's the first time retrieving the TelegramBareClient, the - current authorization is exported to the new DC so that - it can be used there, and the connection is initialized. - - If after using the sender a ConnectionResetError is raised, - this method should be called again with init_connection=True - in order to perform the reconnection. - - If bypass_cache is True, a new client will be exported and - it will not be cached. - """ - # Thanks badoualy/kotlogram on /telegram/api/DefaultTelegramClient.kt - # for clearly showing how to export the authorization! ^^ - - client = self._cached_clients.get(dc_id) - if client and not bypass_cache: - if init_connection: - client.reconnect() - return client - else: - dc = self._get_dc(dc_id) - - # Export the current authorization to the new DC. - export_auth = self(ExportAuthorizationRequest(dc_id)) - - # Create a temporary session for this IP address, which needs - # to be different because each auth_key is unique per DC. - # - # Construct this session with the connection parameters - # (system version, device model...) from the current one. - session = JsonSession(self.session) - session.server_address = dc.ip_address - session.port = dc.port - client = TelegramBareClient( - session, self.api_id, self.api_hash, - timeout=self._timeout - ) - client.connect(exported_auth=export_auth) - - if not bypass_cache: - # Don't go through this expensive process every time. - self._cached_clients[dc_id] = client - return client - def create_new_connection(self, on_dc=None): """Creates a new connection which can be used in parallel with the original TelegramClient. A TelegramBareClient @@ -222,15 +169,10 @@ class TelegramClient(TelegramBareClient): If the client is meant to be used on a different data center, the data center ID should be specified instead. - - Note that TelegramBareClients will not handle automatic - reconnection (i.e. switching to another data center to - download media), and InvalidDCError will be raised in - such case. """ if on_dc is None: - client = TelegramBareClient(self.session, self.api_id, self.api_hash, - proxy=self.proxy) + client = TelegramBareClient( + self.session, self.api_id, self.api_hash, proxy=self.proxy) client.connect() else: client = self._get_exported_client(on_dc, bypass_cache=True) @@ -724,41 +666,6 @@ class TelegramClient(TelegramBareClient): return file_path - def download_file(self, - input_location, - file, - part_size_kb=None, - file_size=None, - progress_callback=None, - on_dc=None): - """Downloads the given InputFileLocation to file (a stream or str). - - If 'progress_callback' is not None, it should be a function that - takes two parameters, (bytes_downloaded, total_bytes). Note that - 'total_bytes' simply equals 'file_size', and may be None. - """ - if on_dc is None: - try: - super().download_file( - input_location, - file, - part_size_kb=part_size_kb, - file_size=file_size, - progress_callback=progress_callback - ) - except FileMigrateError as e: - on_dc = e.new_dc - - if on_dc is not None: - client = self._get_exported_client(on_dc) - client.download_file( - input_location, - file, - part_size_kb=part_size_kb, - file_size=file_size, - progress_callback=progress_callback - ) - # endregion # endregion