Add cancel for upload

-cancel while uploading will delete uploaded files also

Signed-off-by: anas <e.anastayyar@gmail.com>

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'\n<b>User:</b> <a href="tg://user?id={download.message.from_user.id}">{download.message.from_user.first_name}</a> (<code>{download.message.from_user.id}</code>)'
-                if download.status() == MirrorStatus.STATUS_DOWNLOADING or download.status() == MirrorStatus.STATUS_CLONING:
                     msg += f"\n<b>To Stop:</b> <code>/{BotCommands.CancelMirror} {download.gid()}</code>"
                 msg += "\n\n"
                 if STATUS_LIMIT is not None:

bot/helper/mirror_utils/status_utils/upload_status.py

@@ -4,11 +4,12 @@ from bot import DOWNLOAD_DIR
 
 
 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 @@ class UploadStatus(Status):
             return f'{get_readable_time(seconds)}'
         except ZeroDivisionError:
             return '-'
+
+    def gid(self) -> str:
+        return self.__gid
+
+    def download(self):
+        return self.obj

bot/helper/mirror_utils/upload_utils/gdriveTools.py

@@ -51,6 +51,7 @@ class GoogleDriveHelper:
         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 @@ class GoogleDriveHelper:
                             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 @@ class GoogleDriveHelper:
         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 @@ class GoogleDriveHelper:
             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 @@ class GoogleDriveHelper:
                 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 @@ class GoogleDriveHelper:
                 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 @@ class GoogleDriveHelper:
                             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 @@ class GoogleDriveHelper:
                     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'<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>SubFolders: </b><code>{self.total_folders}</code>'
@@ -826,6 +843,7 @@ class GoogleDriveHelper:
                                 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 @@ class GoogleDriveHelper:
             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!')

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

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)
 

bot/modules/mirror.py

@@ -62,6 +62,7 @@ class MirrorListener(listeners.MirrorListeners):
             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 @@ class MirrorListener(listeners.MirrorListeners):
         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 @@ class MirrorListener(listeners.MirrorListeners):
                 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'<a href="tg://user?id={self.message.from_user.id}">{self.message.from_user.first_name}</a>'
+        if uname is not None:
+            men = f'{uname} '
+        sendMessage(men + e_str, self.bot, self.update)
         if count == 0:
             self.clean()
         else: