mirror of
https://github.com/tropicoo/yt-dlp-bot.git
synced 2024-09-20 06:46:08 +08:00
Version 1.2.1
This commit is contained in:
parent
c51883904e
commit
15b2ff9feb
|
@ -2,7 +2,7 @@
|
|||
|
||||
Simple and reliable YouTube Download Telegram Bot.
|
||||
|
||||
Version: 1.2. [Release details](RELEASES.md).
|
||||
Version: 1.2.1. [Release details](RELEASES.md).
|
||||
|
||||
![frames](.assets/download_success.png)
|
||||
|
||||
|
|
18
RELEASES.md
18
RELEASES.md
|
@ -1,3 +1,21 @@
|
|||
## Release 1.2.1
|
||||
|
||||
Release date: March 22, 2023
|
||||
|
||||
## New Features
|
||||
|
||||
N/A (Maintenance release)
|
||||
|
||||
## Important
|
||||
|
||||
1. Disabled SQL logs by default: `SQLALCHEMY_ECHO=False` in `envs/.env_common`
|
||||
|
||||
## Misc
|
||||
|
||||
1. Improved upload log messages - now include file size and whether upload of the file is cached
|
||||
|
||||
---
|
||||
|
||||
## Release 1.2
|
||||
|
||||
Release date: March 13, 2023
|
||||
|
|
|
@ -7,7 +7,7 @@ RUN apk add --no-cache --virtual .build-deps \
|
|||
&& MAKEFLAGS="-j$(nproc)" pip install --no-cache-dir -r requirements.txt \
|
||||
&& apk --purge del .build-deps
|
||||
|
||||
COPY ./app_api ./start.sh ./
|
||||
COPY ./app_api ./start.py ./
|
||||
|
||||
COPY yt_shared /app/yt_shared
|
||||
RUN pip install -e /app/yt_shared
|
||||
|
|
|
@ -8,6 +8,7 @@ def setup_logging() -> None:
|
|||
logging.basicConfig(
|
||||
format=log_format, level=logging.getLevelName(settings.LOG_LEVEL)
|
||||
)
|
||||
|
||||
log = logging.getLogger('sqlalchemy.engine.Engine')
|
||||
log.handlers.pop()
|
||||
try:
|
||||
logging.getLogger('sqlalchemy.engine.Engine').handlers.pop()
|
||||
except IndexError:
|
||||
pass
|
||||
|
|
|
@ -7,7 +7,7 @@ RUN apk add --no-cache --virtual .build-deps \
|
|||
&& MAKEFLAGS="-j$(nproc)" pip install --no-cache-dir -r requirements.txt \
|
||||
&& apk --purge del .build-deps
|
||||
|
||||
COPY ./app_bot ./start.sh ./
|
||||
COPY ./app_bot ./start.py ./
|
||||
|
||||
COPY yt_shared /app/yt_shared
|
||||
RUN pip install -e /app/yt_shared
|
||||
|
|
|
@ -7,15 +7,14 @@ from yt_shared.emoji import SUCCESS_EMOJI
|
|||
from yt_shared.enums import MediaFileType, TaskSource
|
||||
from yt_shared.rabbit.publisher import Publisher
|
||||
from yt_shared.schemas.error import ErrorGeneralPayload
|
||||
from yt_shared.schemas.media import Audio, Video
|
||||
from yt_shared.schemas.media import BaseMedia
|
||||
from yt_shared.schemas.success import SuccessPayload
|
||||
from yt_shared.utils.common import format_bytes
|
||||
from yt_shared.utils.file import remove_dir
|
||||
from yt_shared.utils.tasks.tasks import create_task
|
||||
|
||||
from bot.core.handlers.abstract import AbstractHandler
|
||||
from bot.core.tasks.upload import AudioUploadTask, VideoUploadTask
|
||||
from bot.core.utils import bold, code
|
||||
from bot.core.utils import bold
|
||||
|
||||
|
||||
class SuccessHandler(AbstractHandler):
|
||||
|
@ -54,7 +53,7 @@ class SuccessHandler(AbstractHandler):
|
|||
)
|
||||
await self._publisher.send_download_error(err_payload)
|
||||
|
||||
async def _handle(self, media_object: Audio | Video) -> None:
|
||||
async def _handle(self, media_object: BaseMedia) -> None:
|
||||
try:
|
||||
await self._send_success_text(media_object)
|
||||
if self._upload_is_enabled():
|
||||
|
@ -79,7 +78,7 @@ class SuccessHandler(AbstractHandler):
|
|||
)
|
||||
remove_dir(root_path)
|
||||
|
||||
async def _create_upload_task(self, media_object: Audio | Video) -> None:
|
||||
async def _create_upload_task(self, media_object: BaseMedia) -> None:
|
||||
"""Upload video to Telegram chat."""
|
||||
semaphore = asyncio.Semaphore(value=self._bot.conf.telegram.max_upload_tasks)
|
||||
upload_task_cls = self._UPLOAD_TASK_MAP[media_object.file_type]
|
||||
|
@ -98,14 +97,15 @@ class SuccessHandler(AbstractHandler):
|
|||
exception_message_args=(task_name,),
|
||||
)
|
||||
|
||||
async def _send_success_text(self, media_object: Audio | Video) -> None:
|
||||
text = (
|
||||
f'{SUCCESS_EMOJI} {bold("Downloaded")} {media_object.filename}\n'
|
||||
f'📏 {bold("Size")} {format_bytes(media_object.file_size)}'
|
||||
)
|
||||
@staticmethod
|
||||
def _create_success_text(media_object: BaseMedia) -> str:
|
||||
text = f'{SUCCESS_EMOJI} {bold("Downloaded")} {media_object.filename}'
|
||||
if media_object.saved_to_storage:
|
||||
text = f'{text}\n💾 {bold("Saved to")} {code(media_object.storage_path)}'
|
||||
text = f'{text}\n💾 {bold("Saved to media storage")}'
|
||||
return f'{text}\n📏 {bold("Size")} {media_object.file_size_human()}'
|
||||
|
||||
async def _send_success_text(self, media_object: BaseMedia) -> None:
|
||||
text = self._create_success_text(media_object)
|
||||
for user in self._receiving_users:
|
||||
kwargs = {
|
||||
'chat_id': user.id,
|
||||
|
@ -124,7 +124,7 @@ class SuccessHandler(AbstractHandler):
|
|||
user = self._bot.allowed_users[self._get_sender_id()]
|
||||
return user.upload.upload_video_file
|
||||
|
||||
def _validate_file_size_for_upload(self, media_object: Audio | Video) -> None:
|
||||
def _validate_file_size_for_upload(self, media_object: BaseMedia) -> None:
|
||||
if self._body.context.source is TaskSource.API:
|
||||
max_file_size = self._bot.conf.telegram.api.upload_video_max_file_size
|
||||
else:
|
||||
|
|
|
@ -15,4 +15,7 @@ def setup_logging(suppress_asyncio: bool = True, suppress_urllib3: bool = True)
|
|||
logging.getLogger('urllib3').setLevel(logging.WARNING)
|
||||
|
||||
logging.getLogger('pyrogram').setLevel(logging.WARNING)
|
||||
logging.getLogger('sqlalchemy.engine.Engine').handlers.pop()
|
||||
try:
|
||||
logging.getLogger('sqlalchemy.engine.Engine').handlers.pop()
|
||||
except IndexError:
|
||||
pass
|
||||
|
|
|
@ -3,7 +3,7 @@ import asyncio
|
|||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Coroutine
|
||||
|
||||
from pydantic import StrictFloat, StrictInt, StrictStr
|
||||
from pydantic import StrictBool, StrictFloat, StrictInt, StrictStr
|
||||
from pyrogram.enums import ChatAction, MessageMediaType, ParseMode
|
||||
from pyrogram.types import Animation, Message
|
||||
from pyrogram.types import Audio as _Audio
|
||||
|
@ -13,9 +13,8 @@ from yt_shared.db.session import get_db
|
|||
from yt_shared.repositories.task import TaskRepository
|
||||
from yt_shared.schemas.base import RealBaseModel
|
||||
from yt_shared.schemas.cache import CacheSchema
|
||||
from yt_shared.schemas.media import Audio, Video
|
||||
from yt_shared.schemas.media import BaseMedia, Video
|
||||
from yt_shared.schemas.success import SuccessPayload
|
||||
from yt_shared.utils.common import format_bytes
|
||||
from yt_shared.utils.tasks.abstract import AbstractTask
|
||||
from yt_shared.utils.tasks.tasks import create_task
|
||||
|
||||
|
@ -27,21 +26,22 @@ if TYPE_CHECKING:
|
|||
from bot.core.bot import VideoBot
|
||||
|
||||
|
||||
class AbstractUploadContext(RealBaseModel):
|
||||
class BaseUploadContext(RealBaseModel):
|
||||
caption: StrictStr
|
||||
filename: StrictStr
|
||||
filepath: StrictStr
|
||||
duration: StrictFloat
|
||||
type: MessageMediaType
|
||||
is_cached: StrictBool = False
|
||||
|
||||
|
||||
class VideoUploadContext(AbstractUploadContext):
|
||||
class VideoUploadContext(BaseUploadContext):
|
||||
height: StrictInt
|
||||
width: StrictInt
|
||||
thumb: StrictStr
|
||||
|
||||
|
||||
class AudioUploadContext(AbstractUploadContext):
|
||||
class AudioUploadContext(BaseUploadContext):
|
||||
pass
|
||||
|
||||
|
||||
|
@ -50,7 +50,7 @@ class AbstractUploadTask(AbstractTask, metaclass=abc.ABCMeta):
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
media_object: Audio | Video,
|
||||
media_object: BaseMedia,
|
||||
users: list[BaseUserSchema | UserSchema],
|
||||
bot: 'VideoBot',
|
||||
semaphore: asyncio.Semaphore,
|
||||
|
@ -86,7 +86,7 @@ class AbstractUploadTask(AbstractTask, metaclass=abc.ABCMeta):
|
|||
async def _send_upload_text(self) -> None:
|
||||
text = (
|
||||
f'⬆️ {bold("Uploading")} {self._filename}\n'
|
||||
f'📏 {bold("Size")} {format_bytes(self._media_object.file_size)}'
|
||||
f'📏 {bold("Size")} {self._media_object.file_size_human()}'
|
||||
)
|
||||
coros = []
|
||||
for chat_id in self._upload_chat_ids:
|
||||
|
@ -125,7 +125,13 @@ class AbstractUploadTask(AbstractTask, metaclass=abc.ABCMeta):
|
|||
|
||||
async def _upload_file(self) -> None:
|
||||
for chat_id in chain(self._upload_chat_ids, self._forward_chat_ids):
|
||||
self._log.info('Uploading "%s" to chat id "%d"', self._filename, chat_id)
|
||||
self._log.info(
|
||||
'Uploading "%s" [%s] [cached: %s] to chat id "%d"',
|
||||
self._filename,
|
||||
self._media_object.file_size_human(),
|
||||
self._media_ctx.is_cached,
|
||||
chat_id,
|
||||
)
|
||||
await self._bot.send_chat_action(chat_id, action=self._UPLOAD_ACTION)
|
||||
try:
|
||||
message = await self.__upload(chat_id=chat_id)
|
||||
|
@ -206,23 +212,25 @@ class AudioUploadTask(AbstractUploadTask):
|
|||
self._media_ctx.type = message.media
|
||||
self._media_ctx.filepath = audio.file_id
|
||||
self._media_ctx.duration = audio.duration
|
||||
self._media_ctx.is_cached = True
|
||||
self._cached_message = message
|
||||
|
||||
self._create_cache_task(cache_object=audio)
|
||||
|
||||
def _generate_file_caption(self) -> str:
|
||||
caption_items = [
|
||||
caption_items = (
|
||||
f'{bold("Title:")} {self._media_object.title}',
|
||||
f'{bold("Filename:")} {self._media_object.filename}',
|
||||
f'{bold("Filename:")} {self._filename}',
|
||||
f'{bold("URL:")} {self._ctx.context.url}',
|
||||
f'{bold("Size:")} {format_bytes(self._media_object.file_size)}',
|
||||
]
|
||||
f'{bold("Size:")} {self._media_object.file_size_human()}',
|
||||
)
|
||||
return '\n'.join(caption_items)
|
||||
|
||||
|
||||
class VideoUploadTask(AbstractUploadTask):
|
||||
_UPLOAD_ACTION = ChatAction.UPLOAD_VIDEO
|
||||
_media_ctx: VideoUploadContext
|
||||
_media_object: Video
|
||||
|
||||
def _create_media_context(self) -> VideoUploadContext:
|
||||
return VideoUploadContext(
|
||||
|
@ -246,12 +254,12 @@ class VideoUploadTask(AbstractUploadTask):
|
|||
if caption_conf.include_title:
|
||||
caption_items.append(f'{bold("Title:")} {self._media_object.title}')
|
||||
if caption_conf.include_filename:
|
||||
caption_items.append(f'{bold("Filename:")} {self._media_object.filename}')
|
||||
caption_items.append(f'{bold("Filename:")} {self._filename}')
|
||||
if caption_conf.include_link:
|
||||
caption_items.append(f'{bold("URL:")} {self._ctx.context.url}')
|
||||
if caption_conf.include_size:
|
||||
caption_items.append(
|
||||
f'{bold("Size:")} {format_bytes(self._media_object.file_size)}'
|
||||
f'{bold("Size:")} {self._media_object.file_size_human()}'
|
||||
)
|
||||
return '\n'.join(caption_items)
|
||||
|
||||
|
@ -284,6 +292,7 @@ class VideoUploadTask(AbstractUploadTask):
|
|||
self._media_ctx.type = message.media
|
||||
self._media_ctx.filepath = video.file_id
|
||||
self._media_ctx.thumb = video.thumbs[0].file_id
|
||||
self._media_ctx.is_cached = True
|
||||
self._cached_message = message
|
||||
|
||||
self._create_cache_task(cache_object=video)
|
||||
|
|
|
@ -1 +1 @@
|
|||
__version__ = '1.1'
|
||||
__version__ = '1.2.1'
|
||||
|
|
|
@ -9,7 +9,7 @@ RUN apk add --no-cache --virtual .build-deps \
|
|||
&& MAKEFLAGS="-j$(nproc)" pip install --no-cache-dir -r requirements.txt \
|
||||
&& apk --purge del .build-deps
|
||||
|
||||
COPY ./app_worker ./start.sh ./
|
||||
COPY ./app_worker ./start.py ./
|
||||
|
||||
COPY yt_shared /app/yt_shared
|
||||
RUN pip install -e /app/yt_shared
|
||||
|
|
|
@ -48,7 +48,9 @@ class WorkerLauncher:
|
|||
|
||||
async def _set_yt_dlp_version(self) -> None:
|
||||
curr_version = ytdlp_version.__version__
|
||||
self._log.info('Setting current yt-dlp version (%s)', curr_version)
|
||||
self._log.info(
|
||||
'Saving current yt-dlp version (%s) to the database', curr_version
|
||||
)
|
||||
async for db in get_db():
|
||||
await YtdlpRepository().create_or_update_version(curr_version, db)
|
||||
|
||||
|
|
|
@ -8,6 +8,7 @@ def setup_logging() -> None:
|
|||
logging.basicConfig(
|
||||
format=log_format, level=logging.getLevelName(settings.LOG_LEVEL)
|
||||
)
|
||||
|
||||
log = logging.getLogger('sqlalchemy.engine.Engine')
|
||||
log.handlers.pop()
|
||||
try:
|
||||
logging.getLogger('sqlalchemy.engine.Engine').handlers.pop()
|
||||
except IndexError:
|
||||
pass
|
||||
|
|
|
@ -7,7 +7,12 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
from yt_shared.enums import DownMediaType, TaskStatus
|
||||
from yt_shared.models import Task
|
||||
from yt_shared.repositories.task import TaskRepository
|
||||
from yt_shared.schemas.media import Audio, DownMedia, IncomingMediaPayload, Video
|
||||
from yt_shared.schemas.media import (
|
||||
BaseMedia,
|
||||
DownMedia,
|
||||
IncomingMediaPayload,
|
||||
Video,
|
||||
)
|
||||
from yt_shared.utils.file import remove_dir
|
||||
from yt_shared.utils.tasks.tasks import create_task
|
||||
|
||||
|
@ -153,7 +158,7 @@ class MediaService:
|
|||
video.width = video_streams[0]['width']
|
||||
video.height = video_streams[0]['height']
|
||||
|
||||
def _create_copy_file_task(self, file: Audio | Video) -> asyncio.Task:
|
||||
def _create_copy_file_task(self, file: BaseMedia) -> asyncio.Task:
|
||||
task_name = f'Copy {file.file_type} file to storage task'
|
||||
return create_task(
|
||||
self._copy_file_to_storage(file),
|
||||
|
@ -175,7 +180,7 @@ class MediaService:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
async def _copy_file_to_storage(file: Audio | Video) -> None:
|
||||
async def _copy_file_to_storage(file: BaseMedia) -> None:
|
||||
dst = os.path.join(settings.STORAGE_PATH, file.filename)
|
||||
await asyncio.to_thread(shutil.copy2, file.filepath, dst)
|
||||
file.saved_to_storage = True
|
||||
|
|
|
@ -4,7 +4,6 @@ RUN apk add --no-cache \
|
|||
tzdata \
|
||||
htop \
|
||||
bash \
|
||||
netcat-openbsd \
|
||||
libstdc++
|
||||
|
||||
WORKDIR /app
|
||||
|
|
|
@ -19,7 +19,7 @@ services:
|
|||
restart: unless-stopped
|
||||
ports:
|
||||
- "1984:8000"
|
||||
command: bash -c "./start.sh && uvicorn main:app --host 0.0.0.0 --port 8000"
|
||||
command: bash -c "python start.py && uvicorn main:app --host 0.0.0.0 --port 8000"
|
||||
depends_on:
|
||||
- postgres
|
||||
- rabbitmq
|
||||
|
@ -35,7 +35,7 @@ services:
|
|||
- envs/.env_bot
|
||||
restart: unless-stopped
|
||||
command: >
|
||||
bash -c "./start.sh && python main.py"
|
||||
bash -c "python start.py && python main.py"
|
||||
depends_on:
|
||||
- postgres
|
||||
- rabbitmq
|
||||
|
@ -53,7 +53,7 @@ services:
|
|||
- envs/.env_worker
|
||||
restart: unless-stopped
|
||||
command: >
|
||||
bash -c "alembic upgrade head && ./start.sh && python main.py"
|
||||
bash -c "python start.py && alembic upgrade head && python main.py"
|
||||
depends_on:
|
||||
- postgres
|
||||
- rabbitmq
|
||||
|
|
|
@ -8,7 +8,7 @@ POSTGRES_HOST=yt_postgres
|
|||
POSTGRES_PORT=5432
|
||||
POSTGRES_DB=yt
|
||||
|
||||
SQLALCHEMY_ECHO=True
|
||||
SQLALCHEMY_ECHO=False
|
||||
SQLALCHEMY_EXPIRE_ON_COMMIT=False
|
||||
|
||||
RABBITMQ_USER=guest
|
||||
|
|
106
start.py
Normal file
106
start.py
Normal file
|
@ -0,0 +1,106 @@
|
|||
"""Over-engineered Python 3.10+ version of bash script with netcat (nc) just for fun.
|
||||
|
||||
#!/bin/bash
|
||||
|
||||
check_reachability() {
|
||||
while ! nc -z "$1" "${!2}"
|
||||
do
|
||||
echo "Waiting for $3 to be reachable on port ${!2}"
|
||||
sleep 1
|
||||
done
|
||||
echo "Connection to $3 on port ${!2} verified"
|
||||
return 0
|
||||
}
|
||||
|
||||
|
||||
wait_for_services_to_be_reachable() {
|
||||
check_reachability rabbitmq RABBITMQ_PORT RabbitMQ
|
||||
check_reachability postgres POSTGRES_PORT PostgreSQL
|
||||
}
|
||||
|
||||
wait_for_services_to_be_reachable
|
||||
exit 0
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Type
|
||||
|
||||
SOCK_CONNECTED = 0
|
||||
DEFAULT_PORT = 0
|
||||
DEFAULT_SLEEP_TIME = 1
|
||||
|
||||
|
||||
class ServiceRegistry(type):
|
||||
REGISTRY: dict[str, type['BaseService']] = {}
|
||||
|
||||
def __new__(
|
||||
mcs: Type['ServiceRegistry'],
|
||||
name: str,
|
||||
bases: tuple[type['BaseService']],
|
||||
attrs: dict,
|
||||
) -> type['BaseService']:
|
||||
service_cls: type['BaseService'] = type.__new__(mcs, name, bases, attrs)
|
||||
mcs.REGISTRY[service_cls.__name__] = service_cls
|
||||
return service_cls
|
||||
|
||||
@classmethod
|
||||
def get_services(mcs) -> dict[str, type['BaseService']]:
|
||||
return mcs.REGISTRY.copy()
|
||||
|
||||
|
||||
@dataclass
|
||||
class BaseService:
|
||||
name: str = field(default='', init=False)
|
||||
host: str = field(default='', init=False)
|
||||
port: int = field(default=DEFAULT_PORT, init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.__class__ is BaseService:
|
||||
raise TypeError('Cannot instantiate abstract class.')
|
||||
|
||||
|
||||
@dataclass
|
||||
class RabbitMQService(BaseService, metaclass=ServiceRegistry):
|
||||
name: str = field(default='RabbitMQ')
|
||||
host: str = field(default=os.getenv('RABBITMQ_HOST'))
|
||||
port: int = field(default=int(os.getenv('RABBITMQ_PORT')))
|
||||
|
||||
|
||||
@dataclass
|
||||
class PostgreSQLService(BaseService, metaclass=ServiceRegistry):
|
||||
name: str = field(default='PostgreSQL')
|
||||
host: str = field(default=os.getenv('POSTGRES_HOST'))
|
||||
port: int = field(default=int(os.getenv('POSTGRES_PORT')))
|
||||
|
||||
|
||||
async def is_port_open(host: str, port: int) -> bool:
|
||||
sock: socket.socket
|
||||
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
|
||||
return await asyncio.to_thread(sock.connect_ex, (host, port)) == SOCK_CONNECTED
|
||||
|
||||
|
||||
async def check_reachability(service: BaseService) -> None:
|
||||
while True:
|
||||
print(f'Waiting for {service.name} to be reachable on port {service.port}')
|
||||
if await is_port_open(host=service.host, port=service.port):
|
||||
break
|
||||
await asyncio.sleep(DEFAULT_SLEEP_TIME)
|
||||
print(f'Connection to {service.name} on port {service.port} verified')
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
logging.getLogger('asyncio').setLevel(logging.ERROR)
|
||||
coros = []
|
||||
for service_cls in ServiceRegistry.get_services().values():
|
||||
coros.append(check_reachability(service_cls()))
|
||||
await asyncio.gather(*coros)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(asyncio.run(main()))
|
20
start.sh
20
start.sh
|
@ -1,20 +0,0 @@
|
|||
#!/bin/bash
|
||||
|
||||
check_reachability() {
|
||||
while ! nc -z "$1" "${!2}"
|
||||
do
|
||||
echo "Waiting for $3 to be reachable on port ${!2}"
|
||||
sleep 1
|
||||
done
|
||||
echo "Connection to $3 on port ${!2} verified"
|
||||
return 0
|
||||
}
|
||||
|
||||
|
||||
wait_for_services_to_be_reachable() {
|
||||
check_reachability rabbitmq RABBITMQ_PORT RabbitMQ
|
||||
check_reachability postgres POSTGRES_PORT PostgreSQL
|
||||
}
|
||||
|
||||
wait_for_services_to_be_reachable
|
||||
exit 0
|
|
@ -7,7 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
from yt_shared.enums import TaskStatus
|
||||
from yt_shared.models import Cache, File, Task
|
||||
from yt_shared.schemas.cache import CacheSchema
|
||||
from yt_shared.schemas.media import Audio, IncomingMediaPayload, Video
|
||||
from yt_shared.schemas.media import BaseMedia, IncomingMediaPayload, Video
|
||||
from yt_shared.utils.common import ASYNC_LOCK
|
||||
|
||||
|
||||
|
@ -60,7 +60,7 @@ class TaskRepository:
|
|||
|
||||
@staticmethod
|
||||
async def save_file(
|
||||
db: AsyncSession, task: Task, media: Audio | Video, meta: dict
|
||||
db: AsyncSession, task: Task, media: BaseMedia, meta: dict
|
||||
) -> File:
|
||||
file = File(
|
||||
title=media.title,
|
||||
|
|
|
@ -4,14 +4,19 @@ from yt_shared.enums import RabbitPayloadType
|
|||
|
||||
|
||||
class RealBaseModel(BaseModel):
|
||||
"""Base Pydantic model. All models should inherit from this."""
|
||||
|
||||
class Config:
|
||||
extra = Extra.forbid
|
||||
|
||||
def json(self, *args, **kwargs) -> str:
|
||||
"""By default, dump without whitespaces."""
|
||||
if 'separators' not in kwargs:
|
||||
kwargs['separators'] = (',', ':')
|
||||
return super().json(*args, **kwargs)
|
||||
|
||||
|
||||
class BaseRabbitPayloadModel(RealBaseModel):
|
||||
"""Base RabbitMQ payload model. All RabbitMQ models should inherit from this."""
|
||||
|
||||
type: RabbitPayloadType
|
||||
|
|
|
@ -13,9 +13,12 @@ from pydantic import (
|
|||
|
||||
from yt_shared.enums import DownMediaType, MediaFileType, TaskSource, TelegramChatType
|
||||
from yt_shared.schemas.base import RealBaseModel
|
||||
from yt_shared.utils.common import format_bytes
|
||||
|
||||
|
||||
class IncomingMediaPayload(RealBaseModel):
|
||||
"""RabbitMQ incoming media payload from Telegram Bot or API service."""
|
||||
|
||||
id: uuid.UUID | None
|
||||
from_chat_id: StrictInt | None
|
||||
from_chat_type: TelegramChatType | None
|
||||
|
@ -28,30 +31,38 @@ class IncomingMediaPayload(RealBaseModel):
|
|||
added_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
class CommonMedia(RealBaseModel):
|
||||
class BaseMedia(RealBaseModel):
|
||||
"""Model representing abstract downloaded media with common fields."""
|
||||
|
||||
file_type: MediaFileType
|
||||
title: StrictStr
|
||||
filename: StrictStr
|
||||
filepath: StrictStr
|
||||
file_size: StrictInt
|
||||
duration: StrictFloat | None = None
|
||||
orm_file_id: uuid.UUID | None = None
|
||||
|
||||
saved_to_storage: StrictBool = False
|
||||
storage_path: StrictStr | None = None
|
||||
|
||||
def file_size_human(self) -> str:
|
||||
return format_bytes(num=self.file_size)
|
||||
|
||||
|
||||
class Audio(BaseMedia):
|
||||
"""Model representing downloaded audio file."""
|
||||
|
||||
class Audio(CommonMedia):
|
||||
file_type: Literal[MediaFileType.AUDIO] = MediaFileType.AUDIO
|
||||
orm_file_id: uuid.UUID | None = None
|
||||
|
||||
|
||||
class Video(CommonMedia):
|
||||
class Video(BaseMedia):
|
||||
"""Model representing downloaded video file with separate thumbnail."""
|
||||
|
||||
file_type: Literal[MediaFileType.VIDEO] = MediaFileType.VIDEO
|
||||
thumb_name: StrictStr | None = None
|
||||
width: int | None = None
|
||||
height: int | None = None
|
||||
thumb_path: StrictStr | None = None
|
||||
orm_file_id: uuid.UUID | None = None
|
||||
|
||||
@root_validator(pre=False)
|
||||
def _set_fields(cls, values: dict) -> dict:
|
||||
|
|
|
@ -20,4 +20,5 @@ def remove_dir(dir_path: str) -> None:
|
|||
|
||||
|
||||
def file_size(filepath: str) -> int:
|
||||
"""Return file size in bytes."""
|
||||
return os.path.getsize(filepath)
|
||||
|
|
Loading…
Reference in a new issue