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

Improve gdrive progress tracking 

* Still bugged

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>
This commit is contained in:
lzzy12 2019-10-13 23:00:05 +05:30
parent 7228461ab9
commit a0cdda9e6e
5 changed files with 69 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
bot/helper/ext_utils/bot_utils.py
View File

@ -42,7 +42,7 @@ def get_download_status_list():
def get_progress_bar_string(status):
if status.status() == MirrorStatus.STATUS_UPLOADING:
completed = status.uploaded_bytes / 8
completed = status.upload_helper.uploaded_bytes / 8
else:
completed = status.download().completed_length / 8
total = status.download().total_length / 8

28
bot/helper/mirror_utils/download_status.py
View File

@ -14,8 +14,7 @@ class DownloadStatus:
self.__gid = gid
self.__download = get_download(gid)
self.__uid = message_id
self.uploaded_bytes = 0
self.upload_time = 0
self.upload_helper = None
def __update(self):
self.__download = get_download(self.__gid)
@ -26,12 +25,12 @@ class DownloadStatus:
:return: returns progress in percentage
"""
self.__update()
if self.status() == MirrorStatus.STATUS_UPLOADING:
if self.upload_helper is not None:
return f'{round(self.upload_progress(), 2)}%'
return self.__download.progress_string()
def upload_progress(self):
return self.uploaded_bytes / self.download().total_length * 100
return self.upload_helper.uploaded_bytes / self.download().total_length * 100
def __size(self):
"""
@ -42,17 +41,13 @@ class DownloadStatus:
def __upload_speed(self):
"""
Calculates upload speed in bytes/second
:return: Upload speed in Bytes/Seconds
"""
try:
return self.uploaded_bytes / self.upload_time
except ZeroDivisionError:
return 0
return self.upload_helper.speed()
def speed(self):
self.__update()
if self.status() == MirrorStatus.STATUS_UPLOADING:
if self.upload_helper is not None:
return f'{get_readable_file_size(self.__upload_speed())}/s'
return self.__download.download_speed_string()
@ -69,9 +64,9 @@ class DownloadStatus:
def eta(self):
self.__update()
if self.status() == MirrorStatus.STATUS_UPLOADING:
if self.upload_helper is not None:
try:
seconds = round((self.__size() - self.uploaded_bytes) / self.__upload_speed(), 2)
seconds = (self.__size() - self.upload_helper.uploaded_bytes) / self.__upload_speed()
return f'{get_readable_time(seconds)}'
except ZeroDivisionError:
return '-'
@ -79,20 +74,17 @@ class DownloadStatus:
def status(self):
self.__update()
status = None
if self.is_archiving:
status = MirrorStatus.STATUS_ARCHIVING
if self.__download.is_waiting:
elif self.download().is_waiting:
status = MirrorStatus.STATUS_WAITING
elif self.download().is_paused:
status = MirrorStatus.STATUS_CANCELLED
elif self.__download.is_complete:
# If download exists and is complete the it must be uploading
# otherwise the gid would have been removed from the download_list
elif self.upload_helper is not None:
status = MirrorStatus.STATUS_UPLOADING
elif self.__download.has_failed:
status = MirrorStatus.STATUS_FAILED
elif self.__download.is_active:
else:
status = MirrorStatus.STATUS_DOWNLOADING
return status

66
bot/helper/mirror_utils/gdriveTools.py
View File

@ -25,33 +25,49 @@ class GoogleDriveHelper:
self.__G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download"
self.__listener = listener
self.__service = self.authorize()
self.uploadedBytes = 0
self._file_uploaded_bytes = 0
self.uploaded_bytes = 0
self.start_time = 0
self.total_time = 0
self._should_update = True
self._do_progress_update = True
self.is_uploading = True
self.is_cancelled = False
self.status = None
def cancel(self):
self.is_cancelled = True
self.is_uploading = False
def speed(self):
"""
It calculates the average upload speed and returns it in bytes/seconds unit
:return: Upload speed in bytes/second
"""
try:
return self.uploaded_bytes / self.total_time
except ZeroDivisionError:
return 0
def _on_upload_progress(self):
while self._do_progress_update:
while self.is_uploading:
if self.status is not None:
time_lapsed = time.time() - self.start_time
# Update the message only if status is not null and loop_count is multiple of 50
chunk_size = self.status.total_size * self.status.progress() - self.uploadedBytes
self.uploadedBytes = self.status.total_size * self.status.progress()
# LOGGER.info(f'{file_name}: {status.progress() * 100}')
with download_dict_lock:
download_dict[self.__listener.uid].uploaded_bytes += chunk_size
download_dict[self.__listener.uid].upload_time = time_lapsed
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.info(f'Chunk size: {get_readable_file_size(chunk_size)}')
self.uploaded_bytes += chunk_size
self.total_time += DOWNLOAD_STATUS_UPDATE_INTERVAL
if self._should_update:
try:
LOGGER.info('Updating messages')
_list = get_download_status_list()
index = get_download_index(_list, get_download(self.__listener.message.message_id).gid)
self.__listener.onUploadProgress(_list, index)
except KillThreadException:
except KillThreadException as e:
LOGGER.info(f'Stopped calling onDownloadProgress(): {str(e)}')
self._should_update = False
else:
LOGGER.info('status: None')
time.sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
def upload_file(self, file_path, file_name, mime_type, parent_id):
@ -78,10 +94,12 @@ class GoogleDriveHelper:
# Insert a file
drive_file = self.__service.files().create(body=file_metadata, media_body=media_body)
response = None
threading.Thread(target=self._on_upload_progress).start()
file_start_time = time.time()
while response is None:
if self.is_cancelled:
return None
self.status, response = drive_file.next_chunk()
self._do_progress_update = False
self._file_uploaded_bytes = 0
# Insert new permissions
self.__service.permissions().create(fileId=response['id'], body=permissions).execute()
# Define file instance and get url for download
@ -97,22 +115,28 @@ class GoogleDriveHelper:
file_path = f"{file_dir}/{file_name}"
LOGGER.info("Uploading File: " + file_name)
self.start_time = time.time()
threading.Thread(target=self._on_upload_progress).start()
if os.path.isfile(file_path):
try:
mime_type = get_mime_type(file_path)
g_drive_link = self.upload_file(file_path, file_name, mime_type, parent_id)
link = self.upload_file(file_path, file_name, mime_type, parent_id)
if link is None:
raise Exception('Upload has been manually cancelled')
LOGGER.info("Uploaded To G-Drive: " + file_path)
link = g_drive_link
except Exception as e:
LOGGER.error(str(e))
e_str = str(e).replace('<', '')
e_str = e_str.replace('>', '')
self.__listener.onUploadError(e_str, _list, index)
return
finally:
self.is_uploading = False
else:
try:
dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id)
self.upload_dir(file_path, dir_id)
result = self.upload_dir(file_path, dir_id)
if result is None:
raise Exception('Upload has been manually cancelled!')
LOGGER.info("Uploaded To G-Drive: " + file_name)
link = f"https://drive.google.com/folderview?id={dir_id}"
except Exception as e:
@ -121,6 +145,8 @@ class GoogleDriveHelper:
e_str = e_str.replace('>', '')
self.__listener.onUploadError(e_str, _list, index)
return
finally:
self.is_uploading = False
LOGGER.info(download_dict)
self.__listener.onUploadComplete(link, _list, index)
LOGGER.info("Deleting downloaded file/folder..")
@ -152,6 +178,8 @@ class GoogleDriveHelper:
new_id = None
for item in list_dirs:
current_file_name = os.path.join(input_directory, item)
if self.is_cancelled:
return None
if os.path.isdir(current_file_name):
current_dir_id = self.create_directory(item, parent_id)
new_id = self.upload_dir(current_file_name, current_dir_id)

4
bot/modules/cancel_mirror.py
View File

@ -24,6 +24,10 @@ def cancel_mirror(update: Update, context):
downloads = aria2.get_downloads(download.followed_by_ids)
aria2.pause(downloads)
aria2.pause([download])
with download_dict_lock:
upload_helper = download_dict[mirror_message.message_id].upload_helper
if upload_helper is not None:
upload_helper.cancel()
clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/')

10
bot/modules/mirror.py
View File

@ -1,5 +1,5 @@
from telegram.ext import CommandHandler, run_async
from telegram.error import BadRequest
from telegram.error import BadRequest, TimedOut
from bot.helper.mirror_utils import download_tools, gdriveTools, listeners
from bot import LOGGER, dispatcher, DOWNLOAD_DIR
from bot.helper.ext_utils import fs_utils, bot_utils
@ -49,6 +49,8 @@ class MirrorListener(listeners.MirrorListeners):
download_dict[self.uid].is_archiving = False
download_dict[self.uid].upload_name = name
gdrive = gdriveTools.GoogleDriveHelper(self)
with download_dict_lock:
download_dict[self.uid].upload_helper = gdrive
gdrive.upload(name)
def onDownloadError(self, error, progress_status_list: list, index: int):
@ -109,8 +111,10 @@ class MirrorListener(listeners.MirrorListeners):
msg = get_readable_message(progress)
try:
editMessage(msg, self.context, self.reply_message)
except BadRequest:
raise KillThreadException('Message deleted. Do not call this method from the thread')
except BadRequest as e:
raise KillThreadException(str(e))
except TimedOut:
pass
def _mirror(update, context, isTar=False):