add shorts download

Author: mithmith

.env.example

@@ -21,3 +21,4 @@ TG_ADMIN_ID = 1234567890
 # MONITOR_HISTORY = 1
 # MONITOR_VIDEO_FORMATS = 0
 # RUN_TG_BOT = 1
+# RUN_TG_BOT_SHORTS_PUBLISH = 0

app/db/data_table.py

@@ -92,6 +92,7 @@ class Video(Base, table=True):
     upload_date: Optional[datetime] = Field(default=None)
     defaultaudiolanguage: Optional[str] = Field(default=None)
     last_update: datetime = Field(default_factory=lambda: datetime.now().replace(microsecond=0))
+    tg_post_date: Optional[datetime] = Field(default=None)
 
     channel: Channel = Relationship(back_populates="videos")
     thumbnails: List["Thumbnail"] = Relationship(back_populates="video")

app/db/repository.py

@@ -554,7 +554,7 @@ def update_video(self, video_schema: VideoSchema) -> None:
         else:
             logger.error(f"Video with ID {video_schema.id} not found in the database.")
 
-    def update_video_path(self, video_id: UUID, video_path: Path) -> None:
+    def update_video_path(self, video_id: str, video_path: Path) -> None:
         """
         Updates the file path where the video is stored.
 
@@ -567,7 +567,7 @@ def update_video_path(self, video_id: UUID, video_path: Path) -> None:
             ID exists, its 'video_path' attribute is updated to the new path. The method commits the change
             to the database. If the video does not exist, it logs a warning.
         """
-        video: Video = self._session.query(Video).filter_by(id=video_id).first()
+        video: Video = self._session.query(Video).filter_by(video_id=video_id).first()
         if video:
             video.video_path = str(video_path)
             self._session.commit()
@@ -680,3 +680,15 @@ def bulk_add_tags(self, tags: list[str]) -> None:
             except IntegrityError as e:
                 self._session.rollback()  # Rollback transaction in case of an error
                 logger.error(f"Error during bulk add tags: {e}")
+
+
+class YoutubeVideoRepository(BaseRepository[Video]):
+    model = Video
+
+    def update_tg_post_date(self, video_id: str):
+        """Обновляет поле tg_post_date для видео после успешной отправки в Telegram."""
+        with self._session.begin():
+            video: Video = self._session.query(Video).filter(Video.video_id == video_id).first()
+            if video:
+                video.tg_post_date = datetime.now().replace(microsecond=0)
+                self._session.commit()

app/integrations/ytdlp.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import subprocess
 from pathlib import Path
@@ -11,7 +12,7 @@
 from app.db.base import Session
 from app.db.data_table import Video
 from app.db.repository import YoutubeDataRepository
-from app.schema import ChannelInfoSchema, VideoSchema, YTFormatSchema
+from app.schema import ChannelInfoSchema, VideoDownloadSchema, VideoSchema, YTFormatSchema
 
 
 class YTChannelDownloader:
@@ -81,13 +82,21 @@ def _extract_video_list(self) -> tuple[list[VideoSchema], str]:
                     video_list.append(VideoSchema(**entry))
         return video_list, channel_id
 
-    def download_video(self, video_id: str, format: str = "bv+ba/b") -> None:
-        video: Video = self._repository.get_video(video_id)
-        if video:
-            video_path = self._construct_video_path(video_id)
-            command = f'yt-dlp -f "{format}" -o "{video_path}" {video.url}'
-            subprocess.run(command, shell=True, check=True)
-            self._repository.update_video_path(video_id, video_path)
+    @staticmethod
+    async def download_video(video_info: VideoDownloadSchema, format: str = "best[ext=mp4]") -> None:
+        command = f'yt-dlp -f "{format}" -o "{video_info.video_download_path}" {video_info.video_url}'
+        logger.debug("Downloading video: {}".format(video_info.video_url))
+        logger.debug("Executing command: {}".format(command))
+        process = await asyncio.create_subprocess_shell(
+            command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+        )
+
+        _, stderr = await process.communicate()
+
+        if process.returncode == 0:
+            logger.info(f"Видео скачано: {video_info.video_download_path}")
+        else:
+            logger.error(f"Ошибка скачивания видео: {stderr.decode().strip()}")
 
     def download_thumbnail(self, video_id: str) -> None:
         video: Video = self._repository.get_video(video_id)

app/schema.py

@@ -43,6 +43,14 @@ class NewVideoSchema(BaseModel):
     channel_url: str = ""
     video_title: str = ""
     video_url: str = ""
+    video_id: str = ""
+
+
+class VideoDownloadSchema(BaseModel):
+    file_name: str = ""
+    video_download_path: str = ""
+    video_url: str = ""
+    video_id: str = ""
 
 
 class ChannelInfoSchema(BaseModel):

app/service/telegram.py

@@ -10,6 +10,8 @@
 from telegram.ext import Application
 
 from app.config import settings
+from app.db.base import Session
+from app.db.repository import YoutubeVideoRepository
 from app.integrations.telegram import get_telegram_handlers
 from app.schema import NewVideoSchema
 
@@ -24,6 +26,7 @@ def __init__(self, bot_token: str, group_id: str, queue: Queue, delay: int = 30)
         self._delay = delay  # Задержка между отправками сообщений с новыми видео
         self._max_retries = 3  # Максимальное количество попыток запуска бота и отправки сообщений
         self._retry_delay = 5  # Задержка между неудачными попытками (в секундах)
+        self._repository = YoutubeVideoRepository(session=Session())
         logger.info("Telegram bot is created")
 
     def run(self):
@@ -110,7 +113,9 @@ async def _publish_messages(self, bot: Bot):
                 )
                 logger.debug(f"Sending message to {self._group_id}:\n{message}")
 
-                await self._send_message_with_retries(bot=bot, chat_id=self._group_id, text=message)
+                await self._send_message_with_retries(
+                    bot=bot, chat_id=self._group_id, text=message, video_id=video.video_id
+                )
 
                 # Задержка между отправками сообщений
                 await asyncio.sleep(self._delay)
@@ -121,7 +126,7 @@ async def _publish_messages(self, bot: Bot):
             except Exception as e:
                 logger.error(f"Ошибка при отправке сообщения: {e}")
 
-    async def _send_message_with_retries(self, bot: Bot, chat_id: str, text: str):
+    async def _send_message_with_retries(self, bot: Bot, chat_id: str, text: str, video_id: str):
         """
         Отправляет сообщение в Telegram с заданным числом повторных попыток.
 
@@ -136,6 +141,7 @@ async def _send_message_with_retries(self, bot: Bot, chat_id: str, text: str):
                     text=text,
                     parse_mode="Markdown",
                 )
+                self._repository.update_tg_post_date(video_id)
                 logger.info("Сообщение успешно отправлено")
                 return  # Успешная отправка, выходим из функции
             except TelegramError as te:

app/service/yt_monitor.py

@@ -1,15 +1,18 @@
 import asyncio
 from multiprocessing import Process, Queue
+from os.path import join
+from queue import Empty
 from typing import Optional
 
 from loguru import logger
 
+from app.config import settings
 from app.db.base import Session
 from app.db.data_table import Channel, ChannelHistory, Thumbnail
 from app.db.repository import YoutubeDataRepository
 from app.integrations.ytapi import YTApiClient
 from app.integrations.ytdlp import YTChannelDownloader
-from app.schema import ChannelAPIInfoSchema, ChannelInfoSchema, NewVideoSchema, VideoSchema
+from app.schema import ChannelAPIInfoSchema, ChannelInfoSchema, NewVideoSchema, VideoDownloadSchema, VideoSchema
 
 
 class YTMonitorService:
@@ -28,6 +31,7 @@ def __init__(
         self._history_timeout = history_timeout
         self._queue = new_videos_queue  # Очередь для обработки новых видео
         self._shorts_publish = shorts_publish
+        self._download_queue = Queue()
 
     def run(
         self, monitor_new: bool = True, monitor_history: bool = True, monitor_video_formats: bool = True
@@ -50,6 +54,11 @@ def run(
             processes.append(video_formats_process)
             video_formats_process.start()
 
+        if self._shorts_publish:
+            shorts_publish_process = Process(target=self._start_async_loop, args=(self._shorts_downloader,))
+            processes.append(shorts_publish_process)
+            shorts_publish_process.start()
+
         return processes
 
     def _start_async_loop(self, coro_func, *args, **kwargs):
@@ -96,6 +105,21 @@ async def _update_video_formats(self):
             logger.info(f"(FORMATS) Waiting for {self._history_timeout} seconds")
             await asyncio.sleep(self._history_timeout)
 
+    async def _shorts_downloader(self, delay: int = 5):
+        while True:
+            try:
+                video: VideoDownloadSchema = self._download_queue.get(block=False, timeout=5)
+                logger.debug(f"video={video}")
+                await YTChannelDownloader.download_video(video)
+                # Задержка между скачиванием файлов
+                await asyncio.sleep(delay)
+            except Empty:
+                # Если очередь пуста после таймаута, ничего не делаем и продолжаем ждать
+                await asyncio.sleep(delay)
+                continue
+            except Exception as e:
+                logger.error(f"Ошибка при отправке сообщения: {e}")
+
     async def _process_channel_videos(self, channel_url: str, process_new: bool = False, process_old: bool = False):
         """Обработка новых и старых видео для канала."""
         # Получение информации о канале через yt-dlp
@@ -178,11 +202,26 @@ async def _process_channel_videos(self, channel_url: str, process_new: bool = Fa
                                 channel_url=ytdlp_channel_info.channel_url,
                                 video_title=video.title,
                                 video_url=video.url,
+                                video_id=video.id,
                             )
                         )  # add video to queue for telegram bot
-                    # elif self._shorts_publish:
-                    #     self._queue.put(
-                    #     )
+                    elif self._shorts_publish:
+                        video_file_name = (
+                            (ytdlp_channel_info.original_url or ytdlp_channel_info.channel)
+                            + "_shorts_"
+                            + video.id
+                            + ".mp4"
+                        )
+                        self._download_queue.put(
+                            VideoDownloadSchema(
+                                file_name=video_file_name,
+                                video_download_path=join(
+                                    settings.storage_path, settings.shorts_download_path, video_file_name
+                                ),
+                                video_url=video.url,
+                                video_id=video.id,
+                            )
+                        )
         if process_old and old_videos:
             self._process_old_videos(old_videos)