evanferrao
/
dump-mirrorbot
mirror of https://github.com/Box-boi/dump-mirrorbot.git
1
0
Fork 0
Code Issues Actions Packages Projects Releases Wiki Activity

fix sa switch and auto terminate stuck clone (#248) 

Signed-off-by: anas <e.anastayyar@gmail.com>
This commit is contained in:
Anas 2021-06-29 06:55:38 +03:00 committed by GitHub
parent 1d9950c1bb
commit 26dc55ea84
Signed by untrusted user: GitHub
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 33 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
bot/helper/mirror_utils/upload_utils/gdriveTools.py
View File

@ -64,6 +64,7 @@ class GoogleDriveHelper:
self.total_bytes = 0 self.total_bytes = 0
self.total_files = 0 self.total_files = 0
self.total_folders = 0 self.total_folders = 0
self.sa_count = 0
def cancel(self): def cancel(self):
self.is_cancelled = True self.is_cancelled = True
@ -142,11 +143,11 @@ class GoogleDriveHelper:
global SERVICE_ACCOUNT_INDEX global SERVICE_ACCOUNT_INDEX
service_account_count = len(os.listdir("accounts")) service_account_count = len(os.listdir("accounts"))
if SERVICE_ACCOUNT_INDEX == service_account_count - 1: if SERVICE_ACCOUNT_INDEX == service_account_count - 1:
return False SERVICE_ACCOUNT_INDEX = 0
self.sa_count += 1
SERVICE_ACCOUNT_INDEX += 1 SERVICE_ACCOUNT_INDEX += 1
LOGGER.info(f"Switching to {SERVICE_ACCOUNT_INDEX}.json service account") LOGGER.info(f"Switching to {SERVICE_ACCOUNT_INDEX}.json service account")
self.__service = self.authorize() self.__service = self.authorize()
return True
@retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), @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)) retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG))
@ -200,7 +201,7 @@ class GoogleDriveHelper:
response = None response = None
while response is None: while response is None:
if self.is_cancelled: if self.is_cancelled:
return None break
try: try:
self.status, response = drive_file.next_chunk() self.status, response = drive_file.next_chunk()
except HttpError as err: except HttpError as err:
@ -208,11 +209,11 @@ class GoogleDriveHelper:
reason = json.loads(err.content).get('error').get('errors')[0].get('reason') reason = json.loads(err.content).get('error').get('errors')[0].get('reason')
if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded': if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded':
if USE_SERVICE_ACCOUNTS: if USE_SERVICE_ACCOUNTS:
if not self.switchServiceAccount(): self.switchServiceAccount()
raise err
LOGGER.info(f"Got: {reason}, Trying Again.") LOGGER.info(f"Got: {reason}, Trying Again.")
return self.upload_file(file_path, file_name, mime_type, parent_id) return self.upload_file(file_path, file_name, mime_type, parent_id)
else: else:
self.is_cancelled = True
raise err raise err
else: else:
raise err raise err
@ -295,11 +296,14 @@ class GoogleDriveHelper:
reason = json.loads(err.content).get('error').get('errors')[0].get('reason') reason = json.loads(err.content).get('error').get('errors')[0].get('reason')
if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded': if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded':
if USE_SERVICE_ACCOUNTS: if USE_SERVICE_ACCOUNTS:
if not self.switchServiceAccount(): if self.sa_count > self.service_account_count:
self.is_cancelled = True
raise err raise err
LOGGER.info(f"Got: {reason}, Trying Again.") else:
return self.copyFile(file_id,dest_id) self.switchServiceAccount()
return self.copyFile(file_id,dest_id)
else: else:
self.is_cancelled = True
raise err raise err
else: else:
raise err raise err
@ -335,6 +339,8 @@ class GoogleDriveHelper:
self.transferred_size = 0 self.transferred_size = 0
self.total_files = 0 self.total_files = 0
self.total_folders = 0 self.total_folders = 0
if USE_SERVICE_ACCOUNTS:
self.service_account_count = len(os.listdir("accounts"))
try: try:
file_id = self.getIdFromUrl(link) file_id = self.getIdFromUrl(link)
except (KeyError,IndexError): except (KeyError,IndexError):
@ -346,7 +352,7 @@ class GoogleDriveHelper:
meta = self.getFileMetadata(file_id) meta = self.getFileMetadata(file_id)
if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE: if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE:
dir_id = self.create_directory(meta.get('name'), parent_id) dir_id = self.create_directory(meta.get('name'), parent_id)
result = self.cloneFolder(meta.get('name'), meta.get('name'), meta.get('id'), dir_id) self.cloneFolder(meta.get('name'), meta.get('name'), meta.get('id'), dir_id)
msg += f'<b>Filename: </b><code>{meta.get("name")}</code>\n<b>Size: </b><code>{get_readable_file_size(self.transferred_size)}</code>' msg += f'<b>Filename: </b><code>{meta.get("name")}</code>\n<b>Size: </b><code>{get_readable_file_size(self.transferred_size)}</code>'
msg += f'\n<b>Type: </b><code>Folder</code>' msg += f'\n<b>Type: </b><code>Folder</code>'
msg += f'\n<b>SubFolders: </b><code>{self.total_folders}</code>' msg += f'\n<b>SubFolders: </b><code>{self.total_folders}</code>'
@ -417,8 +423,12 @@ class GoogleDriveHelper:
err = err.last_attempt.exception() err = err.last_attempt.exception()
err = str(err).replace('>', '').replace('<', '') err = str(err).replace('>', '').replace('<', '')
LOGGER.error(err) LOGGER.error(err)
if "User rate limit exceeded." in str(err): if "User rate limit exceeded" in str(err):
msg = "User rate limit exceeded." if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE:
msg += "\nUser rate limit exceeded.\nThis folder is empty or missing files/folders"
return msg, InlineKeyboardMarkup(buttons.build_menu(2))
else:
msg = "User rate limit exceeded."
elif "File not found" in str(err): elif "File not found" in str(err):
msg = "File not found." msg = "File not found."
else: else:
@ -444,17 +454,10 @@ class GoogleDriveHelper:
self.transferred_size += int(file.get('size')) self.transferred_size += int(file.get('size'))
except TypeError: except TypeError:
pass pass
try: self.copyFile(file.get('id'), parent_id)
self.copyFile(file.get('id'), parent_id) new_id = parent_id
new_id = parent_id if self.is_cancelled:
except Exception as e: break
if isinstance(e, RetryError):
LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}")
err = e.last_attempt.exception()
else:
err = e
LOGGER.error(err)
return new_id
@retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), @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)) retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG))
@ -479,8 +482,6 @@ class GoogleDriveHelper:
new_id = None new_id = None
for item in list_dirs: for item in list_dirs:
current_file_name = os.path.join(input_directory, item) current_file_name = os.path.join(input_directory, item)
if self.is_cancelled:
return None
if os.path.isdir(current_file_name): if os.path.isdir(current_file_name):
current_dir_id = self.create_directory(item, parent_id) current_dir_id = self.create_directory(item, parent_id)
new_id = self.upload_dir(current_file_name, current_dir_id) new_id = self.upload_dir(current_file_name, current_dir_id)
@ -492,6 +493,8 @@ class GoogleDriveHelper:
self.upload_file(current_file_name, file_name, mime_type, parent_id) self.upload_file(current_file_name, file_name, mime_type, parent_id)
self.total_files += 1 self.total_files += 1
new_id = parent_id new_id = parent_id
if self.is_cancelled:
break
return new_id return new_id
def authorize(self): def authorize(self):
@ -638,7 +641,6 @@ class GoogleDriveHelper:
else : else :
return '', '' return '', ''
def count(self, link): def count(self, link):
try: try:
file_id = self.getIdFromUrl(link) file_id = self.getIdFromUrl(link)
@ -674,9 +676,6 @@ class GoogleDriveHelper:
except TypeError: except TypeError:
pass pass
except Exception as err: except Exception as err:
if isinstance(err, RetryError):
LOGGER.info(f"Total Attempts: {err.last_attempt.attempt_number}")
err = err.last_attempt.exception()
err = str(err).replace('>', '').replace('<', '') err = str(err).replace('>', '').replace('<', '')
LOGGER.error(err) LOGGER.error(err)
if "File not found" in str(err): if "File not found" in str(err):
@ -727,9 +726,6 @@ class GoogleDriveHelper:
pass pass
clonesize = self.total_bytes clonesize = self.total_bytes
except Exception as err: except Exception as err:
if isinstance(err, RetryError):
LOGGER.info(f"Total Attempts: {err.last_attempt.attempt_number}")
err = err.last_attempt.exception()
err = str(err).replace('>', '').replace('<', '') err = str(err).replace('>', '').replace('<', '')
LOGGER.error(err) LOGGER.error(err)
if "File not found" in str(err): if "File not found" in str(err):
@ -740,7 +736,6 @@ class GoogleDriveHelper:
return "", clonesize, name return "", clonesize, name
def download(self, link): def download(self, link):
self.is_downloading = True
file_id = self.getIdFromUrl(link) file_id = self.getIdFromUrl(link)
if USE_SERVICE_ACCOUNTS: if USE_SERVICE_ACCOUNTS:
self.service_account_count = len(os.listdir("accounts")) self.service_account_count = len(os.listdir("accounts"))
@ -755,9 +750,6 @@ class GoogleDriveHelper:
os.makedirs(path) os.makedirs(path)
self.download_file(file_id, path, meta.get('name'), meta.get('mimeType')) self.download_file(file_id, path, meta.get('name'), meta.get('mimeType'))
except Exception as err: except Exception as err:
if isinstance(err, RetryError):
LOGGER.info(f"Total Attempts: {err.last_attempt.attempt_number}")
err = err.last_attempt.exception()
err = str(err).replace('>', '').replace('<', '') err = str(err).replace('>', '').replace('<', '')
LOGGER.error(err) LOGGER.error(err)
if "downloadQuotaExceeded" in str(err): if "downloadQuotaExceeded" in str(err):
@ -821,19 +813,20 @@ class GoogleDriveHelper:
reason = json.loads(err.content).get('error').get('errors')[0].get('reason') reason = json.loads(err.content).get('error').get('errors')[0].get('reason')
if reason == 'downloadQuotaExceeded' or reason == 'dailyLimitExceeded': if reason == 'downloadQuotaExceeded' or reason == 'dailyLimitExceeded':
if USE_SERVICE_ACCOUNTS: if USE_SERVICE_ACCOUNTS:
if not self.switchServiceAccount(): if self.sa_count > self.service_account_count:
self.is_cancelled = True self.is_cancelled = True
raise err raise err
LOGGER.info(f"Got: {reason}, Trying Again...") else:
return self.download_file(file_id, path, filename, mime_type) self.switchServiceAccount()
LOGGER.info(f"Got: {reason}, Trying Again...")
return self.download_file(file_id, path, filename, mime_type)
else: else:
self.is_cancelled = True self.is_cancelled = True
raise err raise err
else: else:
self.is_cancelled = True
raise err raise err
self._file_downloaded_bytes = 0 self._file_downloaded_bytes = 0
def _on_download_progress(self): def _on_download_progress(self):
if self.dstatus is not None: if self.dstatus is not None:
chunk_size = self.dstatus.total_size * self.dstatus.progress() - self._file_downloaded_bytes chunk_size = self.dstatus.total_size * self.dstatus.progress() - self._file_downloaded_bytes