From 8d1bf5aa924d8fcfe4550c27f5884ad80a4dd2d5 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sat, 21 Mar 2020 10:31:24 +0530 Subject: [PATCH] recursively try again instead of handling switch with backoff - Handling service account switching with exponential backoff implementation limits the number of service accounts to be used. recursively retrying again after switch doesnt effect backoff. Signed-off-by: lzzy12 --- .../mirror_utils/upload_utils/gdriveTools.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 0853837..4f795c2 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -46,6 +46,7 @@ class GoogleDriveHelper: self.updater = None self.name = name self.update_interval = 3 + self.service_account_count = len(os.listdir("accounts")) def cancel(self): self.is_cancelled = True @@ -93,6 +94,14 @@ class GoogleDriveHelper: return self.__service.files().create(supportsTeamDrives=True, body=file_metadata, media_body=media_body).execute() + def switchServiceAccount(self): + global SERVICE_ACCOUNT_INDEX + if SERVICE_ACCOUNT_INDEX == self.service_account_count - 1: + SERVICE_ACCOUNT_INDEX = 0 + SERVICE_ACCOUNT_INDEX += 1 + LOGGER.info(f"Switching to {SERVICE_ACCOUNT_INDEX}.json service account") + self.__service = self.authorize() + @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 __set_permission(self, drive_id): @@ -146,12 +155,13 @@ class GoogleDriveHelper: except HttpError as err: if err.resp.get('content-type', '').startswith('application/json'): reason = json.loads(err.content).get('error').get('errors')[0].get('reason') - if reason == 'userRateLimitExceeded': - global SERVICE_ACCOUNT_INDEX - SERVICE_ACCOUNT_INDEX += 1 - LOGGER.info(f"Switching to {SERVICE_ACCOUNT_INDEX}.json service account") - self.__service = self.authorize() - raise err + if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded': + if USE_SERVICE_ACCOUNTS: + self.switchServiceAccount() + LOGGER.info(f"Got: {reason}, Trying Again.") + self.upload_file(file_path, file_name, mime_type, parent_id) + else: + raise err self._file_uploaded_bytes = 0 # Insert new permissions if not IS_TEAM_DRIVE: