yt-dlp-bot/worker/core/launcher.py

import asyncio
import logging

from yt_dlp import version as ytdlp_version

from core.callbacks import rmq_callbacks as cb
from core.config import settings
from yt_shared.db 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):
        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):
        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())