Cache exported TelegramBareClients instead MtProtoSenders

telethon/telegram_bare_client.py

@@ -1,22 +1,22 @@
 import logging
-import platform
 from datetime import timedelta
 from hashlib import md5
 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, InvalidDCError, FloodWaitError
+from .errors import RPCError, FloodWaitError
 from .network import authenticator, MtProtoSender, TcpTransport
 from .utils import get_appropriated_part_size
 
 # For sending and receiving requests
-from .tl import MTProtoRequest, Session, JsonSession
+from .tl import MTProtoRequest
 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
 
 # Easier access for working with media
 from .tl.functions.upload import (
@@ -56,7 +56,6 @@ class TelegramBareClient:
            Session must always be a Session instance, and an optional proxy
            can also be specified to be used on the connection.
         """
-
         self.session = session
         self.api_id = api_id
         self.api_hash = api_hash
@@ -71,11 +70,15 @@ class TelegramBareClient:
 
     # region Connecting
 
-    def connect(self, device_model, system_version, app_version, lang_code):
+    def connect(self, device_model, system_version, app_version, lang_code,
+                exported_auth=None):
         """Connects to the Telegram servers, executing authentication if
            required. Note that authenticating to the Telegram servers is
            not the same as authenticating the desired user itself, which
            may require a call (or several) to 'sign_in' for the first time.
+
+           If 'exported_auth' is not None, it will be used instead to
+           determine the authorization key for the current session.
         """
         transport = TcpTransport(self.session.server_address,
                                  self.session.port, proxy=self.proxy)
@@ -92,17 +95,28 @@ class TelegramBareClient:
 
             # Now it's time to send an InitConnectionRequest
             # This must always be invoked with the layer we'll be using
-            query = InitConnectionRequest(
+            if exported_auth is None:
+                query = GetConfigRequest()
+            else:
+                query = ImportAuthorizationRequest(
+                    exported_auth.id, exported_auth.bytes)
+
+            request = InitConnectionRequest(
                 api_id=self.api_id,
                 device_model=device_model,
                 system_version=system_version,
                 app_version=app_version,
                 lang_code=lang_code,
-                query=GetConfigRequest())
+                query=query)
 
             result = self.invoke(
                 InvokeWithLayerRequest(
-                    layer=layer, query=query))
+                    layer=layer, query=request))
+
+            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.invoke(GetConfigRequest())
 
             # We're only interested in the DC options,
             # although many other options are available!
@@ -180,7 +194,10 @@ class TelegramBareClient:
         except ConnectionResetError:
             self._logger.info('Server disconnected us. Reconnecting and '
                               'resending request...')
-            self.reconnect()
+
+            # TODO Don't actually use these values
+            import platform
+            self.reconnect(platform.node(), platform.system(), self.__version__, 'en')
             return self.invoke(request, timeout=timeout)
 
         except FloodWaitError:
@@ -265,17 +282,17 @@ class TelegramBareClient:
     # region Downloading media
 
     def download_file(self,
-                          input_location,
+                      input_location,
-                          file_path,
+                      file_path,
-                          part_size_kb=None,
+                      part_size_kb=None,
-                          file_size=None,
+                      file_size=None,
-                          progress_callback=None):
+                      progress_callback=None):
         """Downloads the given InputFileLocation to file_path.
 
            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."""
+           'total_bytes' simply equals 'file_size', and may be None.
-
+        """
         if not part_size_kb:
             if not file_size:
                 part_size_kb = 64  # Reasonable default

telethon/telegram_client.py

@@ -116,8 +116,7 @@ class TelegramClient(TelegramBareClient):
         # 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_senders = {}
+        self._cached_clients = {}
-        self._cached_sessions = {}
         self._updates_thread = None
         self._phone_code_hashes = {}
 
@@ -147,11 +146,10 @@ class TelegramClient(TelegramBareClient):
         super(TelegramClient, self).disconnect()
 
         # Also disconnect all the cached senders
-        for sender in self._cached_senders.values():
+        for sender in self._cached_clients.values():
             sender.disconnect()
 
-        self._cached_senders.clear()
+        self._cached_clients.clear()
-        self._cached_sessions.clear()
 
     def reconnect(self, new_dc=None, *args):
         """Disconnects and connects again (effectively reconnecting).
@@ -173,10 +171,10 @@ class TelegramClient(TelegramBareClient):
 
     # region Working with different Data Centers
 
-    def _get_exported_sender(self, dc_id, init_connection=False):
+    def _get_exported_client(self, dc_id, init_connection=False):
-        """Gets a cached exported MtProtoSender for the desired DC.
+        """Gets a cached exported TelegramBareClient for the desired DC.
 
-           If it's the first time retrieving the MtProtoSender, the
+           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.
 
@@ -186,56 +184,36 @@ class TelegramClient(TelegramBareClient):
         # Thanks badoualy/kotlogram on /telegram/api/DefaultTelegramClient.kt
         # for clearly showing how to export the authorization! ^^
 
-        sender = self._cached_senders.get(dc_id)
+        client = self._cached_clients.get(dc_id)
-        session = self._cached_sessions.get(dc_id)
+        if client:
-
-        if sender and session:
             if init_connection:
-                sender.disconnect()
+                client.reconnect(
-                sender.connect()
+                    device_model=self.device_model,
+                    system_version=self.system_version,
+                    app_version=self.app_version,
+                    lang_code=self.lang_code
+                )
 
-            return sender
+            return client
         else:
             dc = self._get_dc(dc_id)
 
-            # Step 1. Export the current authorization to the new DC.
+            # Export the current authorization to the new DC.
             export_auth = self.invoke(ExportAuthorizationRequest(dc_id))
 
-            # Step 2. Create a transport connected to the new DC.
+            # Create a temporary session for this IP address, which needs
-            #         We also create a temporary session because
+            # to be different because each auth_key is unique per DC.
-            #         it's what will contain the required AuthKey
+            session = JsonSession(None)
-            #         for MtProtoSender to work.
+            session.server_address = dc.ip_address
-            transport = TcpTransport(dc.ip_address, dc.port, proxy=self.proxy)
+            session.port = dc.port
-            session = Session(None)
+            client = TelegramBareClient(session, self.api_id, self.api_hash)
-            session.auth_key, session.time_offset = \
+            client.connect(self.device_model, self.system_version,
-                authenticator.do_authentication(transport)
+                           self.app_version, self.lang_code,
+                           exported_auth=export_auth)
 
-            # Step 3. After authenticating on the new DC,
-            #         we can create the proper MtProtoSender.
-            sender = MtProtoSender(transport, session)
-            sender.connect()
-
-            # InvokeWithLayer(InitConnection(ImportAuthorization(...)))
-            init_connection = InitConnectionRequest(
-                api_id=self.api_id,
-                device_model=platform.node(),
-                system_version=platform.system(),
-                app_version=self.__version__,
-                lang_code='en',
-                query=ImportAuthorizationRequest(
-                    export_auth.id, export_auth.bytes)
-            )
-            query = InvokeWithLayerRequest(layer=layer, query=init_connection)
-
-            sender.send(query)
-            sender.receive(query)
-
-            # Step 4. We're connected and using the desired layer!
             # Don't go through this expensive process every time.
-            self._cached_senders[dc_id] = sender
+            self._cached_clients[dc_id] = client
-            self._cached_sessions[dc_id] = session
+            return client
-
-            return sender
 
     # endregion
 
@@ -302,12 +280,10 @@ class TelegramClient(TelegramBareClient):
            ConnectionResetError will be raised if it occurs a second time.
         """
         try:
-            sender = self._get_exported_sender(
+            client = self._get_exported_client(
                 dc_id, init_connection=reconnect)
 
-            sender.send(request)
+            return client.invoke(request)
-            sender.receive(request)
-            return request.result
 
         except ConnectionResetError:
             if reconnect: