Refinements

This commit is contained in:
Taras Terletskyi 2022-10-13 22:55:18 +03:00
parent 1782a56854
commit b7d1854528
37 changed files with 238 additions and 213 deletions

6
.github/dependabot.yml vendored Normal file
View file

@ -0,0 +1,6 @@
version: 2
updates:
- package-ecosystem: 'pip'
directory: '/'
schedule:
interval: 'daily'

View file

@ -8,27 +8,26 @@ Simple and reliable YouTube Download Telegram Bot.
* Trigger video download by sending link to the Telegram bot or by API call * Trigger video download by sending link to the Telegram bot or by API call
* Upload downloaded videos to Telegram * Upload downloaded videos to Telegram
* Track download tasks in the database or API * Track download tasks in the database or API
* Everything is run in Docker containers
## ⚙ Quick Setup ## ⚙ Quick Setup
1. Create Telegram bot using [BotFather](https://t.me/BotFather) and get your `token`. 1. Create Telegram bot using [BotFather](https://t.me/BotFather) and get your `token`
2. [Get your own Telegram API key](https://my.telegram.org/apps) (`api_id` and `api_hash`) 2. [Get own Telegram API key](https://my.telegram.org/apps) (`api_id` and `api_hash`)
3. Find your Telegram User ID [here](https://stackoverflow.com/questions/32683992/find-out-my-own-user-id-for-sending-a-message-with-telegram-api). 3. [Find your Telegram User ID](https://stackoverflow.com/questions/32683992/find-out-my-own-user-id-for-sending-a-message-with-telegram-api)
4. Copy `bot/config-example.yml` to `bot/config.yml`. 4. Copy `bot/config-example.yml` to `bot/config.yml`
5. Write `token`, `api_id`, `api_hash` and your User ID to `bot/config.yml` by changing respective placeholders. 5. Write `token`, `api_id`, `api_hash` and your User ID to `bot/config.yml` by changing respective placeholders
6. Check the default environment variables in `envs/.env_common` and change if needed. 6. Check the default environment variables in `envs/.env_common` and change if needed
7. Video storage path (`STORAGE_PATH` environment variable) is located in `envs/.env_common` file. 7. Video storage path (`STORAGE_PATH` environment variable) is located in the `envs/.env_worker` file
By default, it's `/filestorage` path inside the container. What you want is to map the real path to this inside the `docker-compose.yml` file for `worker` service e.g. By default, it's `/filestorage` path inside the container. What you want is to map the real path to this inside the `docker-compose.yml` file for `worker` service e.g.
if you're on Windows, next strings mean container path `/filestorage` is mapped to if you're on Windows, next strings mean container path `/filestorage` is mapped to
real `D:/Videos` so your videos will be saved to your `Videos` folder. real `D:/Videos` so your videos will be saved to your `Videos` folder.
```yml
worker:
...
volumes:
- "D:/Videos:/filestorage"
```
8. If you want your downloaded video to be uploaded back to Telegram, set `UPLOAD_VIDEO_FILE` 8. If you want your downloaded video to be uploaded back to Telegram, set `UPLOAD_VIDEO_FILE`
environment variable in `.env_common` file to `True`. environment variable in the `.env_bot` file to `True`
```yml
worker:
...
volumes:
- "D:/Videos:/filestorage"
```
## 🏃 Run ## 🏃 Run
Simple as `docker-compose up -d`. Simple as `docker-compose up -d`.
@ -39,7 +38,6 @@ After pasting video URL bot will send you appropriate message whether it was dow
## Advanced setup ## Advanced setup
1. If you want to change `yt-dlp` download options, go to the `worker/ytdl_opts` 1. If you want to change `yt-dlp` download options, go to the `worker/ytdl_opts`
directory, copy content from `default.py` to `user.py` and modify as you wish directory, copy content from `default.py` to `user.py` and modify as you wish
by checking [official documentation](https://github.com/timethrow/yt-dlp/blob/patch-1/README.md#embedding-yt-dlp). by checking [official documentation](https://github.com/timethrow/yt-dlp/blob/patch-1/README.md#embedding-yt-dlp).

View file

@ -7,12 +7,12 @@ from sqlalchemy.ext.asyncio import AsyncEngine
# this is the Alembic Config object, which provides # this is the Alembic Config object, which provides
# access to the values within the .ini file in use. # access to the values within the .ini file in use.
from yt_shared.config import SQLALCHEMY_DATABASE_URI_ASYNC from yt_shared.config import settings
from yt_shared.db import Base from yt_shared.db import Base
from yt_shared.models import * # noqa from yt_shared.models import * # noqa
config = context.config config = context.config
config.set_main_option("sqlalchemy.url", SQLALCHEMY_DATABASE_URI_ASYNC) config.set_main_option('sqlalchemy.url', settings.SQLALCHEMY_DATABASE_URI_ASYNC)
# Interpret the config file for Python logging. # Interpret the config file for Python logging.
# This line sets up loggers basically. # This line sets up loggers basically.
@ -27,7 +27,7 @@ target_metadata = Base.metadata
# other values from the config, defined by the needs of env.py, # other values from the config, defined by the needs of env.py,
# can be acquired: # can be acquired:
# my_important_option = config.get_main_option("my_important_option") # my_important_option = config.get_main_option('my_important_option')
# ... etc. # ... etc.
@ -43,12 +43,12 @@ def run_migrations_offline():
script output. script output.
""" """
url = config.get_main_option("sqlalchemy.url") url = config.get_main_option('sqlalchemy.url')
context.configure( context.configure(
url=url, url=url,
target_metadata=target_metadata, target_metadata=target_metadata,
literal_binds=True, literal_binds=True,
dialect_opts={"paramstyle": "named"}, dialect_opts={'paramstyle': 'named'},
) )
with context.begin_transaction(): with context.begin_transaction():
@ -72,7 +72,7 @@ async def run_migrations_online():
connectable = AsyncEngine( connectable = AsyncEngine(
engine_from_config( engine_from_config(
config.get_section(config.config_ini_section), config.get_section(config.config_ini_section),
prefix="sqlalchemy.", prefix='sqlalchemy.',
poolclass=pool.NullPool, poolclass=pool.NullPool,
future=True, future=True,
) )

View file

@ -1,21 +1,25 @@
FROM python:3.10-slim-buster FROM python:3.10-alpine
RUN apt update \ RUN apk add --no-cache \
&& apt install --yes --no-install-recommends \ tzdata \
bash htop net-tools iputils-ping procps netcat \ htop \
&& rm -rf /var/lib/apt/lists/* bash \
netcat-openbsd
COPY ./api/requirements.txt ./yt_shared/requirements_shared.txt /app/ COPY ./api/requirements.txt ./yt_shared/requirements_shared.txt /app/
WORKDIR /app WORKDIR /app
RUN apt update \ RUN apk add --no-cache --virtual .build-deps \
&& apt install --yes gcc g++ libffi-dev libjpeg-dev zlib1g-dev python3-dev build-essential libtool automake \ linux-headers \
libffi-dev \
zlib-dev \
build-base \
&& pip install --upgrade pip setuptools wheel \ && pip install --upgrade pip setuptools wheel \
&& MAKEFLAGS="-j$(nproc)" \
&& pip install --no-cache-dir -r requirements.txt -r requirements_shared.txt \ && pip install --no-cache-dir -r requirements.txt -r requirements_shared.txt \
&& apt-get autoremove --yes gcc g++ libffi-dev libjpeg-dev zlib1g-dev python3-dev build-essential libtool automake \ && rm requirements_shared.txt \
&& rm -rf /var/lib/apt/lists/* \ && apk --purge del .build-deps
&& rm requirements_shared.txt
COPY ./api /app COPY ./api /app

View file

@ -1,11 +1,12 @@
import aioredis
from fastapi import FastAPI from fastapi import FastAPI
from fastapi.middleware.gzip import GZipMiddleware from fastapi.middleware.gzip import GZipMiddleware
from fastapi_cache import FastAPICache from fastapi_cache import FastAPICache
from fastapi_cache.backends.redis import RedisBackend from fastapi_cache.backends.redis import RedisBackend
from redis import asyncio as aioredis
from api.api_v1.urls import v1_router from api.api_v1.urls import v1_router
from api.root.endpoints.healthcheck import healthcheck_router from api.root.endpoints.healthcheck import healthcheck_router
from core.config import settings
from core.constants import GZIP_MIN_SIZE from core.constants import GZIP_MIN_SIZE
from yt_shared.rabbit import get_rabbitmq from yt_shared.rabbit import get_rabbitmq
@ -28,7 +29,7 @@ app = create_app()
@app.on_event('startup') @app.on_event('startup')
async def startup_event(): async def startup_event():
redis = aioredis.from_url( redis = aioredis.from_url(
'redis://yt_redis', encoding='utf8', decode_responses=True settings.REDIS_URL, encoding='utf8', decode_responses=True
) )
FastAPICache.init(RedisBackend(redis), prefix='fastapi-cache') FastAPICache.init(RedisBackend(redis), prefix='fastapi-cache')
await get_rabbitmq().register() await get_rabbitmq().register()

8
api/core/config.py Normal file
View file

@ -0,0 +1,8 @@
from yt_shared.config import Settings
class ApiSettings(Settings):
pass
settings = ApiSettings()

View file

@ -1,8 +0,0 @@
from pydantic import BaseSettings
class Settings(BaseSettings):
title: str = 'test app title or what'
settings = Settings()

View file

@ -1,4 +1,3 @@
fastapi-cache2[redis]==0.1.8 fastapi-cache2[redis]==0.1.9
fastapi==0.78.0 fastapi==0.85.0
uvicorn==0.17.6 uvicorn==0.18.3
uvloop==0.16.0

View file

@ -1,21 +1,28 @@
FROM python:3.10-slim-buster FROM python:3.10-alpine
RUN apt update \ RUN apk add --no-cache \
&& apt install --yes --no-install-recommends \ tzdata \
bash htop net-tools iputils-ping procps netcat \ htop \
&& rm -rf /var/lib/apt/lists/* bash \
netcat-openbsd
COPY ./bot/requirements.txt ./yt_shared/requirements_shared.txt /app/ COPY ./bot/requirements.txt ./yt_shared/requirements_shared.txt /app/
WORKDIR /app WORKDIR /app
RUN apt update \ RUN apk add --no-cache --virtual .build-deps \
&& apt install --yes gcc g++ libffi-dev libjpeg-dev zlib1g-dev python3-dev build-essential libtool automake \ linux-headers \
libffi-dev \
zlib-dev \
build-base \
musl-dev \
openssl-dev \
python3-dev \
&& pip install --upgrade pip setuptools wheel \ && pip install --upgrade pip setuptools wheel \
&& MAKEFLAGS="-j$(nproc)" \
&& pip install --no-cache-dir -r requirements.txt -r requirements_shared.txt \ && pip install --no-cache-dir -r requirements.txt -r requirements_shared.txt \
&& apt-get autoremove --yes gcc g++ libffi-dev libjpeg-dev zlib1g-dev python3-dev build-essential libtool automake \ && rm requirements_shared.txt \
&& rm -rf /var/lib/apt/lists/* \ && apk --purge del .build-deps
&& rm requirements_shared.txt
COPY ./bot /app COPY ./bot /app

View file

@ -1,14 +1,11 @@
import asyncio
import logging import logging
from pyrogram import Client from pyrogram import Client
from pyrogram.enums import ParseMode from pyrogram.enums import ParseMode
from core.config.config import get_main_config from core.config.config import get_main_config
from core.tasks.manager import RabbitWorkerManager
from core.tasks.ytdlp import YtdlpNewVersionNotifyTask
from core.utils import bold from core.utils import bold
from yt_shared.rabbit import get_rabbitmq
from yt_shared.task_utils.tasks import create_task
class VideoBot(Client): class VideoBot(Client):
@ -26,20 +23,12 @@ class VideoBot(Client):
self._log.info('Initializing bot client') self._log.info('Initializing bot client')
self.user_ids: list[int] = conf.telegram.allowed_user_ids self.user_ids: list[int] = conf.telegram.allowed_user_ids
self.rabbit_mq = get_rabbitmq()
self.rabbit_worker_manager = RabbitWorkerManager(bot=self)
async def start_tasks(self): @staticmethod
await self.rabbit_worker_manager.start_worker_tasks() async def run_forever() -> None:
"""Firstly 'await bot.start()' should be called."""
task_name = YtdlpNewVersionNotifyTask.__class__.__name__ while True:
create_task( await asyncio.sleep(86400)
YtdlpNewVersionNotifyTask(bot=self).run(),
task_name=task_name,
logger=self._log,
exception_message='Task %s raised an exception',
exception_message_args=(task_name,),
)
async def send_startup_message(self) -> None: async def send_startup_message(self) -> None:
"""Send welcome message after bot launch.""" """Send welcome message after bot launch."""
@ -58,5 +47,5 @@ class VideoBot(Client):
await self.send_message(user_id, text, parse_mode=parse_mode) await self.send_message(user_id, text, parse_mode=parse_mode)
except Exception: except Exception:
self._log.exception( self._log.exception(
'Failed to send message "%s" to user ID ' '%s', text, user_id 'Failed to send message "%s" to user ID %s', text, user_id
) )

View file

@ -1,4 +1,3 @@
import asyncio
import logging import logging
from pyrogram import filters from pyrogram import filters
@ -7,6 +6,7 @@ from pyrogram.handlers import MessageHandler
from core.bot import VideoBot from core.bot import VideoBot
from core.callbacks import TelegramCallback from core.callbacks import TelegramCallback
from core.config.config import get_main_config from core.config.config import get_main_config
from core.tasks.manager import RabbitWorkerManager
from yt_shared.rabbit import get_rabbitmq from yt_shared.rabbit import get_rabbitmq
from yt_shared.task_utils.tasks import create_task from yt_shared.task_utils.tasks import create_task
@ -24,14 +24,15 @@ class BotLauncher:
logging.getLogger().setLevel(get_main_config().log_level) logging.getLogger().setLevel(get_main_config().log_level)
self._bot = VideoBot() self._bot = VideoBot()
self._rabbit_mq = get_rabbitmq() self._rabbit_mq = get_rabbitmq()
self._rabbit_worker_manager = RabbitWorkerManager(bot=self._bot)
async def run(self) -> None: async def run(self) -> None:
"""Run bot.""" """Run bot."""
await self._setup_rabbit() await self._setup_rabbit()
await self._setup_handlers() self._setup_handlers()
await self._start_bot() await self._start_bot()
async def _setup_handlers(self): def _setup_handlers(self):
cb = TelegramCallback() cb = TelegramCallback()
self._bot.add_handler( self._bot.add_handler(
MessageHandler( MessageHandler(
@ -60,16 +61,14 @@ class BotLauncher:
exception_message_args=(task_name,), exception_message_args=(task_name,),
) )
async def _start_tasks(self):
await self._rabbit_worker_manager.start_worker_tasks()
async def _start_bot(self) -> None: async def _start_bot(self) -> None:
"""Start telegram bot and related processes.""" """Start telegram bot and related processes."""
await self._bot.start() await self._bot.start()
self._log.info('Starting %s bot', (await self._bot.get_me()).first_name) self._log.info('Starting %s bot', (await self._bot.get_me()).first_name)
await self._bot.start_tasks() await self._start_tasks()
await self._bot.send_startup_message() await self._bot.send_startup_message()
await self._run_bot_forever() await self._bot.run_forever()
@staticmethod
async def _run_bot_forever() -> None:
while True:
await asyncio.sleep(86400)

View file

@ -0,0 +1,6 @@
from core.config.config import settings
__all__ = [
'settings',
]

View file

@ -10,6 +10,7 @@ from pydantic import ValidationError
from core.config.schema import ConfigSchema from core.config.schema import ConfigSchema
from core.exceptions import ConfigError from core.exceptions import ConfigError
from yt_shared.config import Settings
class ConfigLoader: class ConfigLoader:
@ -58,3 +59,13 @@ _CONF_MAIN = config_loader.load_config()
def get_main_config() -> ConfigSchema: def get_main_config() -> ConfigSchema:
return _CONF_MAIN return _CONF_MAIN
class BotSettings(Settings):
YTDLP_VERSION_CHECK_INTERVAL: int
UPLOAD_VIDEO_FILE: bool
MAX_UPLOAD_VIDEO_FILE_SIZE: int
settings = BotSettings()

View file

@ -5,7 +5,7 @@ from yt_shared.schemas.base import RealBaseModel
_LOG_LEVELS = {'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'} _LOG_LEVELS = {'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'}
_LANG_CODE_LEN = 2 _LANG_CODE_LEN = 2
_LANG_CODE_REGEX = fr'^[a-z]{{{_LANG_CODE_LEN}}}$' _LANG_CODE_REGEX = rf'^[a-z]{{{_LANG_CODE_LEN}}}$'
class TelegramSchema(RealBaseModel): class TelegramSchema(RealBaseModel):

View file

@ -1,5 +1,4 @@
import abc import abc
import asyncio
import enum import enum
import os import os
from typing import Optional, TYPE_CHECKING, Type from typing import Optional, TYPE_CHECKING, Type
@ -8,20 +7,16 @@ from aio_pika import IncomingMessage
from pydantic import BaseModel from pydantic import BaseModel
from pyrogram.enums import ParseMode from pyrogram.enums import ParseMode
from core.config import settings
from core.exceptions import InvalidBodyError from core.exceptions import InvalidBodyError
from core.tasks.abstract import AbstractTask
from core.tasks.upload import UploadTask from core.tasks.upload import UploadTask
from core.utils import bold from core.utils import bold
from yt_shared.config import (
MAX_UPLOAD_VIDEO_FILE_SIZE,
TMP_DOWNLOAD_PATH,
UPLOAD_VIDEO_FILE,
)
from yt_shared.emoji import SUCCESS_EMOJI from yt_shared.emoji import SUCCESS_EMOJI
from yt_shared.rabbit import get_rabbitmq from yt_shared.rabbit import get_rabbitmq
from yt_shared.rabbit.rabbit_config import ERROR_QUEUE, SUCCESS_QUEUE from yt_shared.rabbit.rabbit_config import ERROR_QUEUE, SUCCESS_QUEUE
from yt_shared.schemas.error import ErrorPayload from yt_shared.schemas.error import ErrorPayload
from yt_shared.schemas.success import SuccessPayload from yt_shared.schemas.success import SuccessPayload
from yt_shared.task_utils.abstract import AbstractTask
from yt_shared.task_utils.tasks import create_task from yt_shared.task_utils.tasks import create_task
if TYPE_CHECKING: if TYPE_CHECKING:
@ -88,8 +83,8 @@ class SuccessResultWorkerTask(AbstractResultWorkerTask):
async def _process_body(self, body: SuccessPayload) -> None: async def _process_body(self, body: SuccessPayload) -> None:
await self._send_success_text(body) await self._send_success_text(body)
video_path: str = os.path.join(TMP_DOWNLOAD_PATH, body.filename) video_path: str = os.path.join(settings.TMP_DOWNLOAD_PATH, body.filename)
thumb_path: str = os.path.join(TMP_DOWNLOAD_PATH, body.thumb_name) thumb_path: str = os.path.join(settings.TMP_DOWNLOAD_PATH, body.thumb_name)
if self._eligible_for_upload(video_path): if self._eligible_for_upload(video_path):
await self._create_upload_task(body) await self._create_upload_task(body)
else: else:
@ -106,8 +101,8 @@ class SuccessResultWorkerTask(AbstractResultWorkerTask):
@staticmethod @staticmethod
def _eligible_for_upload(video_path: str) -> bool: def _eligible_for_upload(video_path: str) -> bool:
return ( return (
UPLOAD_VIDEO_FILE settings.UPLOAD_VIDEO_FILE
and os.stat(video_path).st_size <= MAX_UPLOAD_VIDEO_FILE_SIZE and os.stat(video_path).st_size <= settings.MAX_UPLOAD_VIDEO_FILE_SIZE
) )
async def _create_upload_task(self, body: SuccessPayload) -> None: async def _create_upload_task(self, body: SuccessPayload) -> None:
@ -123,12 +118,16 @@ class SuccessResultWorkerTask(AbstractResultWorkerTask):
async def _send_success_text(self, body: SuccessPayload) -> None: async def _send_success_text(self, body: SuccessPayload) -> None:
text = f'{SUCCESS_EMOJI} Downloaded {bold(body.filename)}' text = f'{SUCCESS_EMOJI} Downloaded {bold(body.filename)}'
await self._bot.send_message( kwargs = {
chat_id=body.from_user_id, 'text': text,
reply_to_message_id=body.message_id, 'parse_mode': ParseMode.HTML,
text=text, }
parse_mode=ParseMode.HTML, if body.from_user_id:
) kwargs['chat_id'] = body.from_user_id
if body.message_id:
kwargs['reply_to_message_id'] = body.message_id
await self._bot.send_message(**kwargs)
class ErrorResultWorkerTask(AbstractResultWorkerTask): class ErrorResultWorkerTask(AbstractResultWorkerTask):

View file

@ -5,9 +5,9 @@ from pyrogram.enums import ChatAction
from pyrogram.types import Message from pyrogram.types import Message
from tenacity import retry, stop_after_attempt, wait_fixed from tenacity import retry, stop_after_attempt, wait_fixed
from core.tasks.abstract import AbstractTask from core.config import settings
from yt_shared.task_utils.abstract import AbstractTask
from core.utils import bold from core.utils import bold
from yt_shared.config import TMP_DOWNLOAD_PATH
from yt_shared.db import get_db from yt_shared.db import get_db
from yt_shared.repositories.task import TaskRepository from yt_shared.repositories.task import TaskRepository
from yt_shared.schemas.cache import CacheSchema from yt_shared.schemas.cache import CacheSchema
@ -22,8 +22,8 @@ class UploadTask(AbstractTask):
super().__init__() super().__init__()
self._body = body self._body = body
self.filename = body.filename self.filename = body.filename
self.filepath = os.path.join(TMP_DOWNLOAD_PATH, self.filename) self.filepath = os.path.join(settings.TMP_DOWNLOAD_PATH, self.filename)
self.thumb_path = os.path.join(TMP_DOWNLOAD_PATH, body.thumb_name) self.thumb_path = os.path.join(settings.TMP_DOWNLOAD_PATH, body.thumb_name)
self._bot = bot self._bot = bot
async def run(self) -> None: async def run(self) -> None:
@ -40,8 +40,8 @@ class UploadTask(AbstractTask):
@retry(wait=wait_fixed(10), stop=stop_after_attempt(10), reraise=True) @retry(wait=wait_fixed(10), stop=stop_after_attempt(10), reraise=True)
async def _upload_video_file(self) -> Optional[Message]: async def _upload_video_file(self) -> Optional[Message]:
user_id = self._body.from_user_id user_id = self._body.from_user_id
self._log.info('Uploading to user id %s', user_id)
try: try:
self._log.info('Uploading for user id %s', user_id)
await self._bot.send_chat_action(user_id, action=ChatAction.UPLOAD_VIDEO) await self._bot.send_chat_action(user_id, action=ChatAction.UPLOAD_VIDEO)
return await self._bot.send_video( return await self._bot.send_video(
chat_id=user_id, chat_id=user_id,

View file

@ -2,14 +2,13 @@ import asyncio
import datetime import datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from pyrogram.enums import ParseMode
from core.tasks.abstract import AbstractTask
from core.utils import bold, code from core.utils import bold, code
from yt_shared.config import YTDLP_VERSION_CHECK_INTERVAL
from core.config import settings
from yt_shared.db import get_db from yt_shared.db import get_db
from yt_shared.emoji import INFORMATION_EMOJI from yt_shared.emoji import INFORMATION_EMOJI
from yt_shared.schemas.ytdlp import VersionContext from yt_shared.schemas.ytdlp import VersionContext
from yt_shared.task_utils.abstract import AbstractTask
from yt_shared.ytdlp.version_checker import VersionChecker from yt_shared.ytdlp.version_checker import VersionChecker
if TYPE_CHECKING: if TYPE_CHECKING:
@ -34,13 +33,13 @@ class YtdlpNewVersionNotifyTask(AbstractTask):
'Next yt-dlp version check planned at %s', 'Next yt-dlp version check planned at %s',
self._get_next_check_datetime().isoformat(' '), self._get_next_check_datetime().isoformat(' '),
) )
await asyncio.sleep(YTDLP_VERSION_CHECK_INTERVAL) await asyncio.sleep(settings.YTDLP_VERSION_CHECK_INTERVAL)
@staticmethod @staticmethod
def _get_next_check_datetime() -> datetime.datetime: def _get_next_check_datetime() -> datetime.datetime:
return ( return (
datetime.datetime.now() datetime.datetime.now()
+ datetime.timedelta(seconds=YTDLP_VERSION_CHECK_INTERVAL) + datetime.timedelta(seconds=settings.YTDLP_VERSION_CHECK_INTERVAL)
).replace(microsecond=0) ).replace(microsecond=0)
async def _notify_if_new_version(self) -> None: async def _notify_if_new_version(self) -> None:

View file

@ -1,6 +1,5 @@
PyYAML==6.0 PyYAML==6.0
Pyrogram==2.0.27 Pyrogram==2.0.57
addict==2.4.0 addict==2.4.0
tenacity==8.0.1 tenacity==8.1.0
tgcrypto==1.2.3 tgcrypto==1.2.3
uvloop==0.16.0

View file

@ -1,4 +1,4 @@
version: "3" version: "3.8"
services: services:
api: api:
@ -56,7 +56,7 @@ services:
- "shared-tmpfs:/tmp/download_tmpfs" - "shared-tmpfs:/tmp/download_tmpfs"
postgres: postgres:
container_name: yt_postgres container_name: yt_postgres
image: postgres:14 image: "postgres:14"
env_file: env_file:
- envs/.env_common - envs/.env_common
ports: ports:

View file

@ -1,2 +1,4 @@
APPLICATION_NAME=yt_bot APPLICATION_NAME=yt_bot
YTDLP_VERSION_CHECK_INTERVAL=300 YTDLP_VERSION_CHECK_INTERVAL=300
UPLOAD_VIDEO_FILE=True
MAX_UPLOAD_VIDEO_FILE_SIZE=2147483648

View file

@ -19,7 +19,3 @@ RABBITMQ_PORT=5672
REDIS_HOST=yt_redis REDIS_HOST=yt_redis
TMP_DOWNLOAD_PATH=/tmp/download_tmpfs TMP_DOWNLOAD_PATH=/tmp/download_tmpfs
STORAGE_PATH=/filestorage
SAVE_VIDEO_FILE=True
UPLOAD_VIDEO_FILE=False
MAX_UPLOAD_VIDEO_FILE_SIZE=2147483648

View file

@ -1,2 +1,4 @@
APPLICATION_NAME=yt_worker APPLICATION_NAME=yt_worker
SAVE_VIDEO_FILE=False
MAX_SIMULTANEOUS_DOWNLOADS=10 MAX_SIMULTANEOUS_DOWNLOADS=10
STORAGE_PATH=/filestorage

View file

@ -1,21 +1,26 @@
FROM python:3.10-slim-buster FROM python:3.10-alpine
RUN apt update \ RUN apk add --no-cache \
&& apt install --yes --no-install-recommends \ tzdata \
bash htop net-tools iputils-ping procps netcat ffmpeg \ htop \
&& rm -rf /var/lib/apt/lists/* bash \
netcat-openbsd \
ffmpeg
COPY ./worker/requirements.txt ./yt_shared/requirements_shared.txt /app/ COPY ./worker/requirements.txt ./yt_shared/requirements_shared.txt /app/
WORKDIR /app WORKDIR /app
RUN apt update \ RUN apk add --no-cache --virtual .build-deps \
&& apt install --yes gcc g++ libffi-dev libjpeg-dev zlib1g-dev python3-dev build-essential libtool automake \ linux-headers \
libffi-dev \
zlib-dev \
build-base \
&& pip install --upgrade pip setuptools wheel \ && pip install --upgrade pip setuptools wheel \
&& MAKEFLAGS="-j$(nproc)" \
&& pip install --no-cache-dir -r requirements.txt -r requirements_shared.txt \ && pip install --no-cache-dir -r requirements.txt -r requirements_shared.txt \
&& apt-get autoremove --yes gcc g++ libffi-dev libjpeg-dev zlib1g-dev python3-dev build-essential libtool automake \ && rm requirements_shared.txt \
&& rm -rf /var/lib/apt/lists/* \ && apk --purge del .build-deps
&& rm requirements_shared.txt
COPY ./worker /app COPY ./worker /app
COPY ./alembic /app/alembic COPY ./alembic /app/alembic

11
worker/core/config.py Normal file
View file

@ -0,0 +1,11 @@
from yt_shared.config import Settings
class WorkerSettings(Settings):
APPLICATION_NAME: str
MAX_SIMULTANEOUS_DOWNLOADS: int
SAVE_VIDEO_FILE: bool
STORAGE_PATH: str
settings = WorkerSettings()

View file

@ -4,7 +4,7 @@ import logging
from yt_dlp import version as ytdlp_version from yt_dlp import version as ytdlp_version
from core.callbacks import rmq_callbacks as cb from core.callbacks import rmq_callbacks as cb
from yt_shared.config import MAX_SIMULTANEOUS_DOWNLOADS from core.config import settings
from yt_shared.db import get_db from yt_shared.db import get_db
from yt_shared.rabbit import get_rabbitmq from yt_shared.rabbit import get_rabbitmq
from yt_shared.rabbit.rabbit_config import INPUT_QUEUE from yt_shared.rabbit.rabbit_config import INPUT_QUEUE
@ -26,12 +26,17 @@ class WorkerLauncher:
loop.run_until_complete(self.stop()) loop.run_until_complete(self.stop())
async def _start(self) -> None: 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()) await asyncio.gather(self._setup_rabbit(), self._set_yt_dlp_version())
async def _setup_rabbit(self) -> None: async def _setup_rabbit(self) -> None:
self._log.info('Setting up RabbitMQ connection') self._log.info('Setting up RabbitMQ connection')
await self._rabbit_mq.register() await self._rabbit_mq.register()
await self._rabbit_mq.channel.set_qos(prefetch_count=MAX_SIMULTANEOUS_DOWNLOADS) 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) await self._rabbit_mq.queues[INPUT_QUEUE].consume(cb.on_input_message)
async def _set_yt_dlp_version(self): async def _set_yt_dlp_version(self):

View file

@ -1,19 +1,18 @@
import abc
import asyncio import asyncio
import logging
import os import os
import signal import signal
from typing import Optional from typing import Optional
from yt_shared.task_utils.abstract import AbstractTask
from yt_shared.utils import wrap from yt_shared.utils import wrap
class AbstractFfBinaryTask(metaclass=abc.ABCMeta): class AbstractFfBinaryTask(AbstractTask):
_CMD: Optional[str] = None _CMD: Optional[str] = None
_CMD_TIMEOUT = 10 _CMD_TIMEOUT = 10
def __init__(self, file_path: str) -> None: def __init__(self, file_path: str) -> None:
self._log = logging.getLogger(self.__class__.__name__) super().__init__()
self._file_path = file_path self._file_path = file_path
self._killpg = wrap(os.killpg) self._killpg = wrap(os.killpg)
@ -37,8 +36,3 @@ class AbstractFfBinaryTask(metaclass=abc.ABCMeta):
async def _get_stdout_stderr(proc: asyncio.subprocess.Process) -> tuple[str, str]: async def _get_stdout_stderr(proc: asyncio.subprocess.Process) -> tuple[str, str]:
stdout, stderr = await proc.stdout.read(), await proc.stderr.read() stdout, stderr = await proc.stdout.read(), await proc.stderr.read()
return stdout.decode().strip(), stderr.decode().strip() return stdout.decode().strip(), stderr.decode().strip()
@abc.abstractmethod
async def run(self) -> None:
"""Main entry point."""
pass

View file

@ -6,10 +6,10 @@ from typing import Optional
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from core.config import settings
from core.downloader import VideoDownloader from core.downloader import VideoDownloader
from core.tasks.ffprobe_context import GetFfprobeContextTask from core.tasks.ffprobe_context import GetFfprobeContextTask
from core.tasks.thumbnail import MakeThumbnailTask from core.tasks.thumbnail import MakeThumbnailTask
from yt_shared.config import SAVE_VIDEO_FILE, STORAGE_PATH, TMP_DOWNLOAD_PATH
from yt_shared.constants import TaskStatus from yt_shared.constants import TaskStatus
from yt_shared.models import Task from yt_shared.models import Task
from yt_shared.rabbit.publisher import Publisher from yt_shared.rabbit.publisher import Publisher
@ -39,15 +39,15 @@ class VideoService:
async def _post_process_file( async def _post_process_file(
self, video: DownVideo, task: Task, db: AsyncSession self, video: DownVideo, task: Task, db: AsyncSession
) -> None: ) -> None:
file_path = os.path.join(TMP_DOWNLOAD_PATH, video.name) file_path = os.path.join(settings.TMP_DOWNLOAD_PATH, video.name)
thumb_path = os.path.join(TMP_DOWNLOAD_PATH, video.thumb_name) thumb_path = os.path.join(settings.TMP_DOWNLOAD_PATH, video.thumb_name)
post_process_coros = [ post_process_coros = [
self._set_probe_ctx(file_path, video), self._set_probe_ctx(file_path, video),
MakeThumbnailTask(thumb_path, file_path).run(), MakeThumbnailTask(thumb_path, file_path).run(),
] ]
if SAVE_VIDEO_FILE: if settings.SAVE_VIDEO_FILE:
post_process_coros.append(self._copy_file_to_storage(video)) post_process_coros.append(self._copy_file_to_storage(video))
await asyncio.gather(*post_process_coros) await asyncio.gather(*post_process_coros)
@ -69,8 +69,8 @@ class VideoService:
@staticmethod @staticmethod
async def _copy_file_to_storage(video: DownVideo) -> None: async def _copy_file_to_storage(video: DownVideo) -> None:
src = os.path.join(TMP_DOWNLOAD_PATH, video.name) src = os.path.join(settings.TMP_DOWNLOAD_PATH, video.name)
dst = os.path.join(STORAGE_PATH, video.name) dst = os.path.join(settings.STORAGE_PATH, video.name)
await asyncio.to_thread(shutil.copy2, src, dst) await asyncio.to_thread(shutil.copy2, src, dst)
async def _start_download( async def _start_download(

View file

@ -1,4 +1,4 @@
PyYAML==6.0 PyYAML==6.0
alembic==1.8.0 alembic==1.8.1
python-dotenv==0.20.0 python-dotenv==0.21.0
yt-dlp yt-dlp

View file

@ -1,12 +1,12 @@
import os import os
from yt_shared.config import TMP_DOWNLOAD_PATH from core.config import settings
# More here https://github.com/yt-dlp/yt-dlp/blob/master/yt_dlp/options.py # More here https://github.com/yt-dlp/yt-dlp/blob/master/yt_dlp/options.py
YTDL_OPTS = { YTDL_OPTS = {
'outtmpl': os.path.join(TMP_DOWNLOAD_PATH, '%(title)s.%(ext)s'), 'outtmpl': os.path.join(settings.TMP_DOWNLOAD_PATH, '%(title)s.%(ext)s'),
'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4', 'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4',
'noplaylist': True, 'noplaylist': True,
'max_downloads': 1, 'max_downloads': 1,
'concurrent_fragment_downloads': 2, 'concurrent_fragment_downloads': 5,
} }

View file

@ -1,7 +1,6 @@
SQLAlchemy-Utils==0.38.2 SQLAlchemy-Utils==0.38.3
SQLAlchemy==1.4.37 SQLAlchemy[asyncio]==1.4.41
aio-pika==8.0.3 aio-pika==8.2.2
aiohttp==3.8.1 aiohttp==3.8.3
asyncpg==0.25.0 asyncpg==0.26.0
orjson==3.7.2 pydantic==1.10.2
pydantic==1.9.1

View file

@ -1,40 +1,42 @@
import os from pydantic import BaseSettings, Field
from yt_shared.utils import get_env_bool
APPLICATION_NAME = os.getenv('APPLICATION_NAME', 'APPLICATION_NAME_NOT_SET') class Settings(BaseSettings):
APPLICATION_NAME: str
POSTGRES_USER = os.getenv('POSTGRES_USER', 'yt') POSTGRES_USER: str
POSTGRES_PASSWORD = os.getenv('POSTGRES_PASSWORD', 'yt') POSTGRES_PASSWORD: str
POSTGRES_HOST = os.getenv('POSTGRES_HOST', 'localhost') POSTGRES_HOST: str
POSTGRES_PORT = os.getenv('POSTGRES_PORT', 5432) POSTGRES_PORT: int
POSTGRES_DB = os.getenv('POSTGRES_DB', 'yt') POSTGRES_DB: str
POSTGRES_TEST_DB = os.getenv('POSTGRES_TEST_DB', 'yt_test') POSTGRES_TEST_DB: str = Field(default='yt_test')
SQLALCHEMY_DATABASE_URI_ASYNC = f'postgresql+asyncpg://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB}' @property
SQLALCHEMY_ECHO = get_env_bool(os.getenv('SQLALCHEMY_ECHO', False)) def SQLALCHEMY_DATABASE_URI_ASYNC(self) -> str:
SQLALCHEMY_EXPIRE_ON_COMMIT = get_env_bool( return f'postgresql+asyncpg://{self.POSTGRES_USER}:{self.POSTGRES_PASSWORD}@{self.POSTGRES_HOST}:{self.POSTGRES_PORT}/{self.POSTGRES_DB}'
os.getenv('SQLALCHEMY_EXPIRE_ON_COMMIT', False)
)
RABBITMQ_USER = os.getenv('RABBITMQ_USER', 'guest') SQLALCHEMY_ECHO: bool
RABBITMQ_PASSWORD = os.getenv('RABBITMQ_PASSWORD', 'guest') SQLALCHEMY_EXPIRE_ON_COMMIT: bool
RABBITMQ_HOST = os.getenv('RABBITMQ_HOST', 'localhost')
RABBITMQ_PORT = os.getenv('RABBITMQ_PORT', 5672)
RABBITMQ_URI = (
f'amqp://{RABBITMQ_USER}:{RABBITMQ_PASSWORD}@{RABBITMQ_HOST}:{RABBITMQ_PORT}/'
)
CONSUMER_NUMBER_OF_RETRY = os.getenv('CONSUMER_NUMBER_OF_RETRY', 2) RABBITMQ_USER: str
RESEND_DELAY_MS = os.getenv('RESEND_DELAY_MS', 60000) RABBITMQ_PASSWORD: str
RABBITMQ_HOST: str
RABBITMQ_PORT: int
REDIS_HOST = os.getenv('REDIS_HOST', 'yt_redis') @property
def RABBITMQ_URI(self) -> str:
return f'amqp://{self.RABBITMQ_USER}:{self.RABBITMQ_PASSWORD}@{self.RABBITMQ_HOST}:{self.RABBITMQ_PORT}/'
TMP_DOWNLOAD_PATH = os.getenv('TMP_DOWNLOAD_PATH', '/tmp/download_tmpfs') CONSUMER_NUMBER_OF_RETRY: int = Field(default=2)
STORAGE_PATH = os.getenv('STORAGE_PATH', '/') RESEND_DELAY_MS: int = Field(default=60000)
MAX_SIMULTANEOUS_DOWNLOADS = int(os.getenv('MAX_SIMULTANEOUS_DOWNLOADS', 10))
YTDLP_VERSION_CHECK_INTERVAL = int(os.getenv('YTDLP_VERSION_CHECK_INTERVAL', 86400))
SAVE_VIDEO_FILE = get_env_bool(os.getenv('SAVE_VIDEO_FILE', True)) REDIS_HOST: str
UPLOAD_VIDEO_FILE = get_env_bool(os.getenv('UPLOAD_VIDEO_FILE', False))
MAX_UPLOAD_VIDEO_FILE_SIZE = int(os.getenv('MAX_UPLOAD_VIDEO_FILE_SIZE', 2147483648)) @property
def REDIS_URL(self) -> str:
return f'redis://{self.REDIS_HOST}'
TMP_DOWNLOAD_PATH: str
settings = Settings()

View file

@ -12,18 +12,13 @@ from sqlalchemy.orm import (
) )
from sqlalchemy_utils import UUIDType from sqlalchemy_utils import UUIDType
from yt_shared.config import ( from yt_shared.config import settings
APPLICATION_NAME,
SQLALCHEMY_DATABASE_URI_ASYNC,
SQLALCHEMY_ECHO,
SQLALCHEMY_EXPIRE_ON_COMMIT,
)
engine = create_async_engine( engine = create_async_engine(
SQLALCHEMY_DATABASE_URI_ASYNC, settings.SQLALCHEMY_DATABASE_URI_ASYNC,
echo=SQLALCHEMY_ECHO, echo=settings.SQLALCHEMY_ECHO,
pool_pre_ping=True, pool_pre_ping=True,
connect_args={'server_settings': {'application_name': APPLICATION_NAME}}, connect_args={'server_settings': {'application_name': settings.APPLICATION_NAME}},
) )
metadata = MetaData() metadata = MetaData()
@ -31,7 +26,7 @@ metadata.bind = engine
AsyncSessionLocal = sessionmaker( AsyncSessionLocal = sessionmaker(
engine, engine,
class_=AsyncSession, class_=AsyncSession,
expire_on_commit=SQLALCHEMY_EXPIRE_ON_COMMIT, expire_on_commit=settings.SQLALCHEMY_EXPIRE_ON_COMMIT,
) )

View file

@ -5,7 +5,7 @@ import aio_pika
from aio_pika import RobustChannel, RobustConnection from aio_pika import RobustChannel, RobustConnection
from aio_pika.abc import AbstractRobustExchange, AbstractRobustQueue from aio_pika.abc import AbstractRobustExchange, AbstractRobustQueue
from yt_shared.config import RABBITMQ_URI from yt_shared.config import settings
from yt_shared.rabbit.rabbit_config import get_rabbit_config from yt_shared.rabbit.rabbit_config import get_rabbit_config
@ -30,7 +30,7 @@ class RabbitMQ:
async def _set_connection(self): async def _set_connection(self):
self.connection = await aio_pika.connect_robust( self.connection = await aio_pika.connect_robust(
RABBITMQ_URI, settings.RABBITMQ_URI,
loop=get_running_loop(), loop=get_running_loop(),
reconnect_interval=self.RABBITMQ_RECONNECT_INTERVAL, reconnect_interval=self.RABBITMQ_RECONNECT_INTERVAL,
) )

View file

@ -2,6 +2,7 @@ import logging
from uuid import UUID from uuid import UUID
from sqlalchemy import insert, select from sqlalchemy import insert, select
from sqlalchemy.exc import NoResultFound
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from yt_shared.constants import TaskStatus from yt_shared.constants import TaskStatus
@ -22,7 +23,10 @@ class TaskRepository:
stmt = select(Task).filter_by(id=video_payload.id) stmt = select(Task).filter_by(id=video_payload.id)
task = await db.execute(stmt) task = await db.execute(stmt)
return task.scalar_one() try:
return task.scalar_one()
except NoResultFound:
return await self._create_task(db, video_payload)
@staticmethod @staticmethod
async def _create_task(db: AsyncSession, video_payload: VideoPayload) -> Task: async def _create_task(db: AsyncSession, video_payload: VideoPayload) -> Task:

View file

@ -1,13 +1,6 @@
import orjson
from pydantic import BaseModel, Extra from pydantic import BaseModel, Extra
def orjson_dumps(v, *, default) -> str:
return orjson.dumps(v, default=default).decode()
class RealBaseModel(BaseModel): class RealBaseModel(BaseModel):
class Config: class Config:
extra = Extra.forbid extra = Extra.forbid
json_loads = orjson.loads
json_dumps = orjson_dumps

View file

@ -8,4 +8,5 @@ class AbstractTask(metaclass=abc.ABCMeta):
@abc.abstractmethod @abc.abstractmethod
async def run(self) -> None: async def run(self) -> None:
"""Main entry point."""
pass pass

View file

@ -1,7 +1,6 @@
import asyncio import asyncio
import functools import functools
import logging import logging
from functools import partial, wraps
from typing import Any, Awaitable, Tuple, TypeVar from typing import Any, Awaitable, Tuple, TypeVar