diff --git a/ytdlbot/db.py b/ytdlbot/db.py index 93ef1e5..517ff88 100644 --- a/ytdlbot/db.py +++ b/ytdlbot/db.py @@ -172,6 +172,7 @@ class MySQL: user_id bigint not null, resolution varchar(128) null, method varchar(64) null, + mode varchar(32) default 'Celery' null, constraint settings_pk primary key (user_id) ); diff --git a/ytdlbot/migration.sql b/ytdlbot/migration.sql new file mode 100644 index 0000000..15fc385 --- /dev/null +++ b/ytdlbot/migration.sql @@ -0,0 +1,3 @@ +alter table settings + add mode varchar(32) default 'Celery' null; + diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py index aca24e8..ad5c1cd 100644 --- a/ytdlbot/tasks.py +++ b/ytdlbot/tasks.py @@ -24,6 +24,7 @@ from celery import Celery from celery.worker.control import Panel from pyrogram import idle from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup +from requests_toolbelt.multipart.encoder import MultipartEncoder from client_init import create_app from config import (ARCHIVE_ID, AUDIO_FORMAT, BROKER, ENABLE_CELERY, @@ -124,8 +125,8 @@ def ytdl_download_entrance(bot_msg, client, url): chat_id = bot_msg.chat.id if forward_video(chat_id, url, client): return - - if ENABLE_CELERY: + mode = get_user_settings(str(chat_id))[-1] + if ENABLE_CELERY and mode in [None, "Celery"]: ytdl_download_task.delay(chat_id, bot_msg.message_id, url) else: ytdl_normal_download(bot_msg, client, url) @@ -231,14 +232,14 @@ def get_dl_source(): return "" -def upload_transfer_sh(video_paths) -> "str": - file = {} - for p in video_paths: - file[p.name] = p.open("rb") +def upload_transfer_sh(paths: list) -> "str": + d = {p.name: (p.name, p.open("rb")) for p in paths} + m = MultipartEncoder(fields=d) + headers = {'Content-Type': m.content_type} try: - req = requests.post("https://transfer.sh", files=file) - return req.text - except requests.exceptions as e: + req = requests.post("https://transfer.sh", data=m, headers=headers) + return re.sub(r"https://", "\nhttps://", req.text) + except requests.exceptions.RequestException as e: return f"Upload failed!❌\n\n```{e}```" @@ -264,7 +265,7 @@ def ytdl_normal_download(bot_msg, client, url): bot_msg.edit_text('Download complete. Sending now...') for video_path in video_paths: # normally there's only one video in that path... - filename = pathlib.Path(video_path).name + filename = video_path.name remain = bot_text.remaining_quota_caption(chat_id) st_size = os.stat(video_path).st_size size = sizeof_fmt(st_size) diff --git a/ytdlbot/utils.py b/ytdlbot/utils.py index 3451a15..a9130e2 100644 --- a/ytdlbot/utils.py +++ b/ytdlbot/utils.py @@ -46,7 +46,7 @@ def get_user_settings(user_id: "str") -> "tuple": cur.execute("SELECT * FROM settings WHERE user_id = %s", (user_id,)) data = cur.fetchone() if data is None: - return 100, "high", "video" + return 100, "high", "video", "Celery" return data diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py index 5d8702d..b2d5761 100644 --- a/ytdlbot/ytdl_bot.py +++ b/ytdlbot/ytdl_bot.py @@ -199,6 +199,15 @@ def direct_handler(client: "Client", message: "types.Message"): def settings_handler(client: "Client", message: "types.Message"): chat_id = message.chat.id client.send_chat_action(chat_id, "typing") + data = get_user_settings(str(chat_id)) + set_mode = (data[-1]) + text = {"Local": "Celery", "Celery": "Local"}.get(set_mode, "Local") + mode_text = f"Download mode: **{set_mode}**" + if message.chat.username == OWNER: + extra = [InlineKeyboardButton(f"Change download mode to {text}", callback_data=text)] + else: + extra = [] + markup = InlineKeyboardMarkup( [ [ # First row @@ -211,11 +220,11 @@ def settings_handler(client: "Client", message: "types.Message"): InlineKeyboardButton("Medium Quality", callback_data="medium"), InlineKeyboardButton("Low Quality", callback_data="low"), ], + extra ] ) - data = get_user_settings(str(chat_id)) - client.send_message(chat_id, bot_text.settings.format(data[1], data[2]), reply_markup=markup) + client.send_message(chat_id, bot_text.settings.format(data[1], data[2]) + mode_text, reply_markup=markup) @app.on_message(filters.command(["vip"])) @@ -303,6 +312,13 @@ def audio_callback(client: "Client", callback_query: types.CallbackQuery): audio_entrance(msg, client) +@app.on_callback_query(filters.regex(r"Local|Celery")) +def owner_local_callback(client: "Client", callback_query: types.CallbackQuery): + chat_id = callback_query.message.chat.id + set_user_settings(chat_id, "mode", callback_query.data) + callback_query.answer(f"Download mode was changed to {callback_query.data}") + + def periodic_sub_check(): vip = VIP() for cid, uids in vip.group_subscriber().items():