mirror of
https://github.com/tropicoo/yt-dlp-bot.git
synced 2025-03-03 18:23:34 +08:00
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
import asyncio
|
|
import logging
|
|
|
|
from core.callbacks import rmq_callbacks as cb
|
|
from core.config import settings
|
|
from yt_dlp import version as ytdlp_version
|
|
|
|
from yt_shared.db.session import get_db
|
|
from yt_shared.rabbit import get_rabbitmq
|
|
from yt_shared.rabbit.rabbit_config import INPUT_QUEUE
|
|
from yt_shared.repositories.ytdlp import YtdlpRepository
|
|
|
|
|
|
class WorkerLauncher:
|
|
def __init__(self) -> None:
|
|
self._log = logging.getLogger(self.__class__.__name__)
|
|
self._rabbit_mq = get_rabbitmq()
|
|
|
|
def start(self) -> None:
|
|
self._log.info('Starting download worker instance')
|
|
loop = asyncio.get_event_loop()
|
|
loop.run_until_complete(self._start())
|
|
try:
|
|
loop.run_forever()
|
|
finally:
|
|
loop.run_until_complete(self.stop())
|
|
|
|
async def _start(self) -> None:
|
|
await self._perform_setup()
|
|
|
|
async def _perform_setup(self) -> None:
|
|
await asyncio.gather(*(self._setup_rabbit(), self._set_yt_dlp_version()))
|
|
|
|
async def _setup_rabbit(self) -> None:
|
|
self._log.info('Setting up RabbitMQ connection')
|
|
await self._rabbit_mq.register()
|
|
await self._rabbit_mq.channel.set_qos(
|
|
prefetch_count=settings.MAX_SIMULTANEOUS_DOWNLOADS
|
|
)
|
|
await self._rabbit_mq.queues[INPUT_QUEUE].consume(cb.on_input_message)
|
|
|
|
async def _set_yt_dlp_version(self) -> None:
|
|
curr_version = ytdlp_version.__version__
|
|
self._log.info('Setting current yt-dlp version (%s)', curr_version)
|
|
async for db in get_db():
|
|
await YtdlpRepository().create_or_update_version(curr_version, db)
|
|
|
|
async def stop(self) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
loop.run_until_complete(self._rabbit_mq.close())
|