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: