From 783967571eb42aa725d3eb56f66605279f143c4f Mon Sep 17 00:00:00 2001 From: Pawarit Thongaek Date: Wed, 8 May 2024 22:33:48 +0700 Subject: [PATCH] feat: create task payload add custom_options --- README.md | 8 +++++++- app_api/api/api/api_v1/schemas/task.py | 2 +- app_api/api/core/services/task.py | 1 + app_bot/bot/core/service.py | 1 + app_worker/worker/core/downloader.py | 7 +++++++ yt_shared/yt_shared/schemas/media.py | 1 + 6 files changed, 18 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 22c7075..195b949 100644 --- a/README.md +++ b/README.md @@ -175,7 +175,13 @@ documentations lives at `http://127.0.0.1:1984/docs`. "download_media_type": "AUDIO_VIDEO", "save_to_storage": false, "custom_filename": "cool.mp4", - "automatic_extension": false + "automatic_extension": false, + "custom_options": { + "http_headers": { + "User-Agent": "MyApp/1.0.0", + "Authorization": "Bearer " + } + } } ``` Response diff --git a/app_api/api/api/api_v1/schemas/task.py b/app_api/api/api/api_v1/schemas/task.py index 2ca6f1f..4e62ee6 100644 --- a/app_api/api/api/api_v1/schemas/task.py +++ b/app_api/api/api/api_v1/schemas/task.py @@ -59,7 +59,7 @@ class CreateTaskIn(StrictRealBaseModel): save_to_storage: bool = ... custom_filename: str = ... automatic_extension: bool = ... - + custom_options: dict | None = Field(None, strict=False) class CreateTaskOut(StrictRealBaseModel): id: uuid.UUID diff --git a/app_api/api/core/services/task.py b/app_api/api/core/services/task.py index 84a72cb..c476a4a 100644 --- a/app_api/api/core/services/task.py +++ b/app_api/api/core/services/task.py @@ -81,6 +81,7 @@ class TaskService: ack_message_id=None, custom_filename=task.custom_filename, automatic_extension=task.automatic_extension, + custom_options=task.custom_options, ) if not await publisher.send_for_download(payload): raise TaskServiceError('Failed to create task') diff --git a/app_bot/bot/core/service.py b/app_bot/bot/core/service.py index 6a7576b..551e57e 100644 --- a/app_bot/bot/core/service.py +++ b/app_bot/bot/core/service.py @@ -37,6 +37,7 @@ class UrlService: download_media_type=url.download_media_type, custom_filename=None, automatic_extension=False, + custom_options=None, ) is_sent = await self._rmq_publisher.send_for_download(payload) if not is_sent: diff --git a/app_worker/worker/core/downloader.py b/app_worker/worker/core/downloader.py index 3c447f3..c15bc18 100644 --- a/app_worker/worker/core/downloader.py +++ b/app_worker/worker/core/downloader.py @@ -60,6 +60,13 @@ class MediaDownloader: media_type=media_type, curr_tmp_dir=curr_tmp_dir ) + if media_payload.custom_options: + for key, value in media_payload.custom_options.items(): + try: + ytdl_opts_model.ytdl_opts[key].update(value) + except KeyError: + ytdl_opts_model.ytdl_opts[key] = value + with yt_dlp.YoutubeDL(ytdl_opts_model.ytdl_opts) as ytdl: self._log.info('Downloading %s', url) self._log.info('Downloading to %s', curr_tmp_dir) diff --git a/yt_shared/yt_shared/schemas/media.py b/yt_shared/yt_shared/schemas/media.py index 47da8a0..8908571 100644 --- a/yt_shared/yt_shared/schemas/media.py +++ b/yt_shared/yt_shared/schemas/media.py @@ -38,6 +38,7 @@ class InbMediaPayload(StrictRealBaseModel): download_media_type: DownMediaType custom_filename: str | None automatic_extension: bool + custom_options: dict[str, dict] | None added_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))