diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py index 9a8f1407..423d10de 100644 --- a/telethon/telegram_bare_client.py +++ b/telethon/telegram_bare_client.py @@ -109,22 +109,6 @@ class TelegramBareClient: self.dc_options = result.dc_options return True - except InvalidDCError as e: - # If an InvalidDCError occurs on CONNECTION then it for sure - # means that the user is associated with a different DC, - # which is the only case in which we need to reconnect. - self._logger.info('DC error on initial connection, ' - 'attempting to reconnect at DC {}' - .format(e.new_dc)) - - self.disconnect() - self.session.auth_key = None # Force creating new auth_key - dc = self._get_dc(e.new_dc) - self.session.server_address = dc.ip_address - self.session.port = dc.port - self.session.save() - self.connect(device_model, system_version, app_version, lang_code) - except (RPCError, ConnectionError) as error: # Probably errors from the previous session, ignore them self.disconnect() @@ -138,12 +122,23 @@ class TelegramBareClient: self.sender.disconnect() self.sender = None - def reconnect(self): - """Disconnects and connects again (effectively reconnecting)""" + def reconnect(self, device_model, system_version, app_version, lang_code, + new_dc=None): + """Disconnects and connects again (effectively reconnecting). + + If 'new_dc' is not None, the current authorization key is + removed, the DC used is switched, and a new connection is made. + """ self.disconnect() - # TODO Don't use these parameters - self.connect( - platform.node(), platform.system(), self.__version__, 'en') + + if new_dc is not None: + self.session.auth_key = None # Force creating new auth_key + dc = self._get_dc(new_dc) + self.session.server_address = dc.ip_address + self.session.port = dc.port + self.session.save() + + self.connect(device_model, system_version, app_version, lang_code) # endregion diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py index 19ee48dc..18839553 100644 --- a/telethon/telegram_client.py +++ b/telethon/telegram_client.py @@ -125,7 +125,8 @@ class TelegramClient(TelegramBareClient): # region Connecting - def connect(self): + def connect(self, a=None, b=None, c=None, d=None): + # TODO Use **kwargs instead 4 dummy parameters return super(TelegramClient, self).connect( device_model=self.device_model, system_version=self.system_version, @@ -145,6 +146,15 @@ class TelegramClient(TelegramBareClient): self._cached_senders.clear() self._cached_sessions.clear() + def reconnect(self, new_dc=None, a=None, b=None, c=None, d=None): + # TODO Use **kwargs instead 4 dummy parameters + super(TelegramClient, self).reconnect( + device_model=self.device_model, + system_version=self.system_version, + app_version=self.app_version, + lang_code=self.lang_code, + new_dc=new_dc) + # endregion # region Working with different Data Centers @@ -248,11 +258,21 @@ class TelegramClient(TelegramBareClient): return result except InvalidDCError as e: - self._logger.info('DC error when invoking request, ' - 'attempting to send it on DC {}' - .format(e.new_dc)) + if (e.message.startswith('PHONE_MIGRATE_') or + e.message.startswith('USER_MIGRATE_')): + self._logger.info('DC error when invoking request, ' + 'attempting to reconnect at DC {}' + .format(e.new_dc)) - return self.invoke_on_dc(request, e.new_dc, timeout=timeout) + self.reconnect(new_dc=e.new_dc) + return self.invoke(request, timeout=timeout) + + else: + self._logger.info('DC error when invoking request, ' + 'attempting to send it on DC {}' + .format(e.new_dc)) + + return self.invoke_on_dc(request, e.new_dc, timeout=timeout) finally: self._lock.release()