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

Improved upload progress implementation 

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>
This commit is contained in:
lzzy12 2020-05-21 12:26:58 +05:30
parent cdefe12348
commit 0f7b79f7a4
1 changed files with 18 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -39,21 +39,18 @@ class GoogleDriveHelper:
self.__service = self.authorize() self.__service = self.authorize()
self.__listener = listener self.__listener = listener
self._file_uploaded_bytes = 0 self._file_uploaded_bytes = 0
self.uploaded_bytes = 0 self.__uploaded_bytes = 0
self.UPDATE_INTERVAL = 5
self.start_time = 0 self.start_time = 0
self.total_time = 0
self._should_update = True
self.is_uploading = True
self.is_cancelled = False self.is_cancelled = False
self.status = None
self.updater = None
self.name = name self.name = name
self.update_interval = 3 self.transferred_size = 0
def cancel(self): def cancel(self):
self.is_cancelled = True self.is_cancelled = True
self.is_uploading = False
@property
def uploaded_bytes(self):
return self.__uploaded_bytes
def speed(self): def speed(self):
""" """
@ -61,7 +58,7 @@ class GoogleDriveHelper:
:return: Upload speed in bytes/second :return: Upload speed in bytes/second
""" """
try: try:
return self.uploaded_bytes / self.total_time return self.uploaded_bytes / (time.time() - self.start_time)
except ZeroDivisionError: except ZeroDivisionError:
return 0 return 0
@ -69,7 +66,7 @@ class GoogleDriveHelper:
def getIdFromUrl(link: str): def getIdFromUrl(link: str):
if "folders" in link or "file" in link: if "folders" in link or "file" in link:
regex = r"https://drive\.google\.com/(drive)?/?u?/?\d?/?(mobile)?/?(file)?(folders)?/?d?/([-\w]+)[?+]?/?(w+)?" regex = r"https://drive\.google\.com/(drive)?/?u?/?\d?/?(mobile)?/?(file)?(folders)?/?d?/([-\w]+)[?+]?/?(w+)?"
res = re.search(regex,link) res = re.search(regex, link)
if res is None: if res is None:
raise IndexError("GDrive ID not found.") raise IndexError("GDrive ID not found.")
return res.group(5) return res.group(5)
@ -78,14 +75,6 @@ class GoogleDriveHelper:
@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))
def _on_upload_progress(self):
if self.status is not None:
chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes
self._file_uploaded_bytes = self.status.total_size * self.status.progress()
LOGGER.debug(f'Uploading {self.name}, chunk size: {get_readable_file_size(chunk_size)}')
self.uploaded_bytes += chunk_size
self.total_time += self.update_interval
def __upload_empty_file(self, path, file_name, mime_type, parent_id=None): def __upload_empty_file(self, path, file_name, mime_type, parent_id=None):
media_body = MediaFileUpload(path, media_body = MediaFileUpload(path,
mimetype=mime_type, mimetype=mime_type,
@ -155,12 +144,14 @@ class GoogleDriveHelper:
drive_file = self.__service.files().create(supportsTeamDrives=True, drive_file = self.__service.files().create(supportsTeamDrives=True,
body=file_metadata, media_body=media_body) body=file_metadata, media_body=media_body)
response = None response = None
last_uploaded = 0
while response is None: while response is None:
if self.is_cancelled: if self.is_cancelled:
return None return None
try: try:
self.status, response = drive_file.next_chunk() status, response = drive_file.next_chunk()
except HttpError as err: except HttpError as err:
status = None
if err.resp.get('content-type', '').startswith('application/json'): if err.resp.get('content-type', '').startswith('application/json'):
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':
@ -170,6 +161,10 @@ class GoogleDriveHelper:
self.upload_file(file_path, file_name, mime_type, parent_id) self.upload_file(file_path, file_name, mime_type, parent_id)
else: else:
raise err raise err
if status is not None:
chunk_size = status.total_size * status.progress() - last_uploaded
last_uploaded = status.total_size * status.progress()
self.__uploaded_bytes += chunk_size
self._file_uploaded_bytes = 0 self._file_uploaded_bytes = 0
# Insert new permissions # Insert new permissions
if not IS_TEAM_DRIVE: if not IS_TEAM_DRIVE:
@ -180,14 +175,11 @@ class GoogleDriveHelper:
return download_url return download_url
def upload(self, file_name: str): def upload(self, file_name: str):
if USE_SERVICE_ACCOUNTS:
self.service_account_count = len(os.listdir("accounts"))
self.__listener.onUploadStarted() self.__listener.onUploadStarted()
file_dir = f"{DOWNLOAD_DIR}{self.__listener.message.message_id}" file_dir = f"{DOWNLOAD_DIR}{self.__listener.message.message_id}"
file_path = f"{file_dir}/{file_name}" file_path = f"{file_dir}/{file_name}"
LOGGER.info("Uploading File: " + file_path) LOGGER.info("Uploading File: " + file_path)
self.start_time = time.time() self.start_time = time.time()
self.updater = setInterval(self.update_interval, self._on_upload_progress)
if os.path.isfile(file_path): if os.path.isfile(file_path):
try: try:
mime_type = get_mime_type(file_path) mime_type = get_mime_type(file_path)
@ -204,8 +196,6 @@ class GoogleDriveHelper:
LOGGER.error(err) LOGGER.error(err)
self.__listener.onUploadError(str(err)) self.__listener.onUploadError(str(err))
return return
finally:
self.updater.cancel()
else: else:
try: try:
dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id) dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id)
@ -223,8 +213,6 @@ class GoogleDriveHelper:
LOGGER.error(err) LOGGER.error(err)
self.__listener.onUploadError(str(err)) self.__listener.onUploadError(str(err))
return return
finally:
self.updater.cancel()
LOGGER.info(download_dict) LOGGER.info(download_dict)
self.__listener.onUploadComplete(link) self.__listener.onUploadComplete(link)
LOGGER.info("Deleting downloaded file/folder..") LOGGER.info("Deleting downloaded file/folder..")
@ -238,7 +226,7 @@ class GoogleDriveHelper:
} }
try: try:
res = self.__service.files().copy(supportsAllDrives=True,fileId=file_id,body=body).execute() res = self.__service.files().copy(supportsAllDrives=True, fileId=file_id, body=body).execute()
return res return res
except HttpError as err: except HttpError as err:
if err.resp.get('content-type', '').startswith('application/json'): if err.resp.get('content-type', '').startswith('application/json'):
@ -247,7 +235,7 @@ class GoogleDriveHelper:
if USE_SERVICE_ACCOUNTS: if USE_SERVICE_ACCOUNTS:
self.switchServiceAccount() self.switchServiceAccount()
LOGGER.info(f"Got: {reason}, Trying Again.") LOGGER.info(f"Got: {reason}, Trying Again.")
self.copyFile(file_id,dest_id) self.copyFile(file_id, dest_id)
else: else:
raise err raise err
@ -255,7 +243,7 @@ class GoogleDriveHelper:
self.transferred_size = 0 self.transferred_size = 0
try: try:
file_id = self.getIdFromUrl(link) file_id = self.getIdFromUrl(link)
except (KeyError,IndexError): except (KeyError, IndexError):
msg = "Google drive ID could not be found in the provided link" msg = "Google drive ID could not be found in the provided link"
return msg return msg
msg = "" msg = ""