From dee36f66bad7b1cee21438f2a816982ab6e7a618 Mon Sep 17 00:00:00 2001 From: anas Date: Tue, 20 Jul 2021 13:22:23 +0300 Subject: [PATCH] Add cancel for upload -cancel while uploading will delete uploaded files also Signed-off-by: anas --- bot/helper/ext_utils/bot_utils.py | 8 ++---- .../status_utils/upload_status.py | 9 +++++- .../mirror_utils/upload_utils/gdriveTools.py | 27 ++++++++++++++++-- bot/modules/cancel_mirror.py | 12 ++------ bot/modules/clone.py | 28 +++++++++---------- bot/modules/mirror.py | 11 ++++++-- 6 files changed, 60 insertions(+), 35 deletions(-) diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index a7ec1c3b801..193472f65ab 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -71,8 +71,7 @@ def getDownloadByGid(gid): with download_dict_lock: for dl in download_dict.values(): status = dl.status() - if status != MirrorStatus.STATUS_UPLOADING and status != MirrorStatus.STATUS_ARCHIVING \ - and status != MirrorStatus.STATUS_EXTRACTING: + if status != MirrorStatus.STATUS_ARCHIVING and status != MirrorStatus.STATUS_EXTRACTING: if dl.gid() == gid: return dl return None @@ -80,10 +79,10 @@ def getDownloadByGid(gid): def getAllDownload(): with download_dict_lock: for dlDetails in list(download_dict.values()): - if dlDetails.status() == MirrorStatus.STATUS_DOWNLOADING \ - or dlDetails.status() == MirrorStatus.STATUS_WAITING: + if dlDetails.status() == MirrorStatus.STATUS_DOWNLOADING or dlDetails.status() == MirrorStatus.STATUS_WAITING: if dlDetails: return dlDetails + return None def get_progress_bar_string(status): completed = status.processed_bytes() / 8 @@ -137,7 +136,6 @@ def get_readable_message(): except: pass msg += f'\nUser: {download.message.from_user.first_name} ({download.message.from_user.id})' - if download.status() == MirrorStatus.STATUS_DOWNLOADING or download.status() == MirrorStatus.STATUS_CLONING: msg += f"\nTo Stop: /{BotCommands.CancelMirror} {download.gid()}" msg += "\n\n" if STATUS_LIMIT is not None: diff --git a/bot/helper/mirror_utils/status_utils/upload_status.py b/bot/helper/mirror_utils/status_utils/upload_status.py index 203e660c0e7..46be352a7cf 100644 --- a/bot/helper/mirror_utils/status_utils/upload_status.py +++ b/bot/helper/mirror_utils/status_utils/upload_status.py @@ -4,11 +4,12 @@ class UploadStatus(Status): - def __init__(self, obj, size, listener): + def __init__(self, obj, size, gid, listener): self.obj = obj self.__size = size self.uid = listener.uid self.message = listener.message + self.__gid = gid def path(self): return f"{DOWNLOAD_DIR}{self.uid}" @@ -52,3 +53,9 @@ def eta(self): return f'{get_readable_time(seconds)}' except ZeroDivisionError: return '-' + + def gid(self) -> str: + return self.__gid + + def download(self): + return self.obj diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 34d22cd5587..0a78bac6950 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -51,6 +51,7 @@ def __init__(self, name=None, listener=None): self.start_time = 0 self.total_time = 0 self.dtotal_time = 0 + self.is_uploading = False self.is_downloading = False self.is_cloning = False self.is_cancelled = False @@ -217,9 +218,12 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): return self.upload_file(file_path, file_name, mime_type, parent_id) else: self.is_cancelled = True + LOGGER.info(f"Got: {reason}") raise err else: raise err + if self.is_cancelled: + return self._file_uploaded_bytes = 0 # Insert new permissions if not IS_TEAM_DRIVE: @@ -230,6 +234,8 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): return download_url def upload(self, file_name: str): + self.is_downloading = False + self.is_uploading = True if USE_SERVICE_ACCOUNTS: self.service_account_count = len(os.listdir("accounts")) self.__listener.onUploadStarted() @@ -242,6 +248,8 @@ def upload(self, file_name: str): try: mime_type = get_mime_type(file_path) link = self.upload_file(file_path, file_name, mime_type, parent_id) + if self.is_cancelled: + return if link is None: raise Exception('Upload has been manually cancelled') LOGGER.info("Uploaded To G-Drive: " + file_path) @@ -256,14 +264,21 @@ def upload(self, file_name: str): return finally: self.updater.cancel() + if self.is_cancelled: + return else: try: dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_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}" + if self.is_cancelled: + LOGGER.info("Deleting uploaded data from drive...") + msg = self.deletefile(link) + LOGGER.info(f"{msg}") + return + LOGGER.info("Uploaded To G-Drive: " + file_name) except Exception as e: if isinstance(e, RetryError): LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") @@ -275,11 +290,12 @@ def upload(self, file_name: str): return finally: self.updater.cancel() + if self.is_cancelled: + return files = self.total_files folders = self.total_folders typ = self.typee self.__listener.onUploadComplete(link, size, files, folders, typ) - LOGGER.info("Deleting downloaded file/folder..") return link @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), @@ -305,6 +321,7 @@ def copyFile(self, file_id, dest_id): return self.copyFile(file_id,dest_id) else: self.is_cancelled = True + LOGGER.info(f"Got: {reason}") raise err else: raise err @@ -360,7 +377,7 @@ def clone(self, link): LOGGER.info("Deleting cloned data from drive...") msg = self.deletefile(durl) LOGGER.info(f"{msg}") - return "Your clone has been stopped and cloned data has been deleted!", "cancelled" + return "your clone has been stopped and cloned data has been deleted!", "cancelled" msg += f'Filename: {meta.get("name")}\nSize: {get_readable_file_size(self.transferred_size)}' msg += f'\nType: Folder' msg += f'\nSubFolders: {self.total_folders}' @@ -826,6 +843,7 @@ def download_file(self, file_id, path, filename, mime_type): return self.download_file(file_id, path, filename, mime_type) else: self.is_cancelled = True + LOGGER.info(f"Got: {reason}") raise err else: raise err @@ -845,3 +863,6 @@ def cancel_download(self): self.__listener.onDownloadError('Download stopped by user!') elif self.is_cloning: LOGGER.info(f"Cancelling Clone: {self.name}") + elif self.is_uploading: + LOGGER.info(f"Cancelling upload: {self.name}") + self.__listener.onUploadError('your upload has been stopped and uploaded data has been deleted!') diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py index 6385ac96930..bf70fa0bb05 100644 --- a/bot/modules/cancel_mirror.py +++ b/bot/modules/cancel_mirror.py @@ -42,19 +42,14 @@ def cancel_mirror(update, context): elif not mirror_message: sendMessage(msg, context.bot, update) return - if dl.status() == "Uploading...📤": - sendMessage("Upload in Progress, You Can't Cancel It.", context.bot, update) - return - elif dl.status() == "Archiving...🔐": + if dl.status() == "Archiving...🔐": sendMessage("Archival in Progress, You Can't Cancel It.", context.bot, update) - return elif dl.status() == "Extracting...📂": sendMessage("Extract in Progress, You Can't Cancel It.", context.bot, update) - return else: dl.download().cancel_download() - sleep(3) # incase of any error with ondownloaderror listener, clean_download will delete the folder but the download will stuck in status msg. - clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/') + sleep(3) # incase of any error with ondownloaderror listener + clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/') def cancel_all(update, context): @@ -68,7 +63,6 @@ def cancel_all(update, context): else: gid = dl.gid() dl.download().cancel_download() - sleep(0.5) count += 1 else: break diff --git a/bot/modules/clone.py b/bot/modules/clone.py index 7d0f0decaf3..e2e5d5e0af4 100644 --- a/bot/modules/clone.py +++ b/bot/modules/clone.py @@ -54,6 +54,18 @@ def cloneNode(update, context): Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages)) sendStatusMessage(update, context.bot) result, button = drive.clone(link) + with download_dict_lock: + del download_dict[update.message.message_id] + count = len(download_dict) + try: + if count == 0: + Interval[0].cancel() + del Interval[0] + delete_all_messages() + else: + update_all_messages() + except IndexError: + pass if update.message.from_user.username: uname = f'@{update.message.from_user.username}' else: @@ -61,24 +73,10 @@ def cloneNode(update, context): if uname is not None: cc = f'\n\ncc: {uname}' men = f'{uname} ' - if button == "cancelled": - sendMessage(men + result, context.bot, update) - elif button == "": + if button == "cancelled" or button == "": sendMessage(men + result, context.bot, update) else: sendMarkup(result + cc, context.bot, update, button) - try: - with download_dict_lock: - del download_dict[update.message.message_id] - count = len(download_dict) - if count == 0: - Interval[0].cancel() - del Interval[0] - delete_all_messages() - else: - update_all_messages() - except (IndexError, KeyError): - pass else: sendMessage('Provide G-Drive Shareable Link to Clone.', context.bot, update) diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 14b59dd8b5a..d812caef6a7 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -62,6 +62,7 @@ def onDownloadComplete(self): LOGGER.info(f"Download completed: {download_dict[self.uid].name()}") download = download_dict[self.uid] name = download.name() + gid = download.gid() size = download.size_raw() if name is None: # when pyrogram's media.file_name is of NoneType name = os.listdir(f'{DOWNLOAD_DIR}{self.uid}')[0] @@ -112,7 +113,7 @@ def onDownloadComplete(self): LOGGER.info(f"Upload Name : {up_name}") drive = gdriveTools.GoogleDriveHelper(up_name, self) size = fs_utils.get_path_size(up_path) - upload_status = UploadStatus(drive, size, self) + upload_status = UploadStatus(drive, size, gid, self) with download_dict_lock: download_dict[self.uid] = upload_status update_all_messages() @@ -218,7 +219,13 @@ def onUploadError(self, error): pass del download_dict[self.message.message_id] count = len(download_dict) - sendMessage(e_str, self.bot, self.update) + if self.message.from_user.username: + uname = f"@{self.message.from_user.username}" + else: + uname = f'{self.message.from_user.first_name}' + if uname is not None: + men = f'{uname} ' + sendMessage(men + e_str, self.bot, self.update) if count == 0: self.clean() else: