Optimize Clone implementation

lzzy12 · Jun 23, 2020 · e4ecc60 · e4ecc60
1 parent 0e5c5ca
commit e4ecc60
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 21 deletions.
diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py
@@ -6,6 +6,7 @@
 import re
 import json
 import requests
+import logging
 
 from google.auth.transport.requests import Request
 from google.oauth2 import service_account
@@ -16,7 +17,7 @@
 from tenacity import *
 
 from bot import parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, \
-    USE_SERVICE_ACCOUNTS
+    USE_SERVICE_ACCOUNTS, download_dict
 from bot.helper.ext_utils.bot_utils import *
 from bot.helper.ext_utils.fs_utils import get_mime_type
 
@@ -167,7 +168,7 @@ def upload_file(self, file_path, file_name, mime_type, parent_id):
                         if USE_SERVICE_ACCOUNTS:
                             self.switchServiceAccount()
                             LOGGER.info(f"Got: {reason}, Trying Again.")
-                            self.upload_file(file_path, file_name, mime_type, parent_id)
+                            return self.upload_file(file_path, file_name, mime_type, parent_id)
                     else:
                         raise err
         self._file_uploaded_bytes = 0
@@ -247,10 +248,37 @@ def copyFile(self, file_id, dest_id):
                     if USE_SERVICE_ACCOUNTS:
                         self.switchServiceAccount()
                         LOGGER.info(f"Got: {reason}, Trying Again.")
-                        self.copyFile(file_id,dest_id)
+                        return self.copyFile(file_id,dest_id)
                 else:
                     raise err
 
+    @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5),
+           retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG))
+    def getFileMetadata(self,file_id):
+        return self.__service.files().get(supportsAllDrives=True, fileId=file_id,
+                                              fields="name,id,mimeType,size").execute()
+
+    @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5),
+           retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG))
+    def getFilesByFolderId(self,folder_id):
+        page_token = None
+        q = f"'{folder_id}' in parents"
+        files = []
+        while True:
+            response = self.__service.files().list(supportsTeamDrives=True,
+                                                   includeTeamDriveItems=True,
+                                                   q=q,
+                                                   spaces='drive',
+                                                   pageSize=200,
+                                                   fields='nextPageToken, files(id, name, mimeType,size)',
+                                                   pageToken=page_token).execute()
+            for file in response.get('files', []):
+                files.append(file)
+            page_token = response.get('nextPageToken', None)
+            if page_token is None:
+                break
+        return files
+
     def clone(self, link):
         self.transferred_size = 0
         try:
@@ -304,23 +332,9 @@ def clone(self, link):
         return msg
 
     def cloneFolder(self, name, local_path, folder_id, parent_id):
-        page_token = None
-        q = f"'{folder_id}' in parents"
-        files = []
         LOGGER.info(f"Syncing: {local_path}")
+        files = self.getFilesByFolderId(folder_id)
         new_id = None
-        while True:
-            response = self.__service.files().list(supportsTeamDrives=True,
-                                                   includeTeamDriveItems=True,
-                                                   q=q,
-                                                   spaces='drive',
-                                                   fields='nextPageToken, files(id, name, mimeType,size)',
-                                                   pageToken=page_token).execute()
-            for file in response.get('files', []):
-                files.append(file)
-            page_token = response.get('nextPageToken', None)
-            if page_token is None:
-                break
         if len(files) == 0:
             return parent_id
         for file in files:

diff --git a/bot/modules/clone.py b/bot/modules/clone.py
@@ -1,12 +1,13 @@
-from telegram.ext import CommandHandler, run_async
+from telegram.ext import CommandHandler
 from bot.helper.mirror_utils.upload_utils.gdriveTools import GoogleDriveHelper
 from bot.helper.telegram_helper.message_utils import *
 from bot.helper.telegram_helper.filters import CustomFilters
 from bot.helper.telegram_helper.bot_commands import BotCommands
+from bot.helper.ext_utils.bot_utils import new_thread
 from bot import dispatcher
 
 
-@run_async
+@new_thread
 def cloneNode(update,context):
     args = update.message.text.split(" ",maxsplit=1)
     if len(args) > 1:
@@ -17,7 +18,7 @@ def cloneNode(update,context):
         deleteMessage(context.bot,msg)
         sendMessage(result,context.bot,update)
     else:
-        sendMessage("Provide G-Drive Shareable Link to Clone.",bot,update)
+        sendMessage("Provide G-Drive Shareable Link to Clone.",context.bot,update)
 
 clone_handler = CommandHandler(BotCommands.CloneCommand,cloneNode,filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
 dispatcher.add_handler(clone_handler)