refactoring for new arch, added Redis, fixed filename, added video exists check

This commit is contained in:
2023-08-27 16:27:28 +03:00
committed by nikili0n
parent a2001a0346
commit a336a7b635
9 changed files with 147 additions and 21 deletions

12
src/core/async_queue.py Normal file
View File

@ -0,0 +1,12 @@
import asyncio
from src.core.redis_client import RedisClient
redis = RedisClient()
class AsyncQueue(asyncio.Queue):
async def put(self, item):
await redis.set_task_to_queue(item)
return await super().put(item)

View File

@ -1,10 +1,15 @@
import asyncio
import concurrent.futures as pool
import os.path
import subprocess
from functools import partial
from fastapi import HTTPException
from src.core.async_queue import AsyncQueue
from src.core.rabbitmq import get_messages
from src.core.redis_client import RedisClient
from src.core.result import Result, ResultTypeEnum
from src.core.ydl import VideoDownloader
from src.exceptions.download_exceptions import SiteNotImplementedException
@ -16,12 +21,16 @@ class MasterService:
self.MAX_EXECUTOR_WORKERS = 8
self.executor = pool.ProcessPoolExecutor(max_workers=self.MAX_EXECUTOR_WORKERS,
initializer=executor_initializer)
self.queue = asyncio.Queue()
self.queue = AsyncQueue()
self.rabbit_consumer = get_messages
self.currently_underway = {} # contains currently in progress videos
async def run(self):
# TODO add cleanup procedure for existing exector procs
subprocess.run(
"for pid in $(ps -ef | grep video_downloader_executor_process | awk '{print $2}'); do kill -9 $pid; done",
shell=True, capture_output=True
)
tasks = [self.loop.create_task(self.create_workers()) for i in range(self.MAX_EXECUTOR_WORKERS + 1)]
await asyncio.gather(self.rabbit_consumer(self.queue), *tasks)
@ -29,15 +38,17 @@ class MasterService:
async def create_workers(self):
while True:
video_params = await self.queue.get()
redis = RedisClient()
await redis.del_task_from_queue_and_add_to_tasks(task=video_params)
self.currently_underway[video_params['link']] = video_params
# TODO save current copy of self.currently_underway to DB\file\etc
download_task = self.loop.run_in_executor(self.executor, partial(
MasterService.video_processing_executor, video_params=video_params
)
)
)
if download_task.done():
result = download_task.result()
# process result
result = await download_task
await redis.del_task_from_tasks_and_add_to_task_done(task=video_params)
# TODO process result
self.queue.task_done()
@staticmethod
@ -46,9 +57,16 @@ class MasterService:
"format": video_params["format"],
"merge_output_format": video_params["merge_output_format"],
'outtmpl': video_params["outtmpl"],
"quiet": True
}
downloader = VideoDownloader(link=video_params["link"], ydl_opts=ydl_opts)
video_info = downloader.get_info()
if os.path.exists(
os.path.join(os.getcwd() + f"Youtube/{video_info['id']}_{video_info['width']}.{video_info['ext']}")
):
return Result(result_type=ResultTypeEnum.EXIST)
try:
downloader.ydl_opts["quiet"] = False
result = downloader.download()
return result
except SiteNotImplementedException as ex:
@ -59,7 +77,6 @@ class MasterService:
@staticmethod
def video_processing_executor(video_params: dict):
# TODO check if video exists on server etc. Result(type=ResultType.EXISTS, value=False)
try:
MasterService.video_download(video_params=video_params)
return Result(result_type=ResultTypeEnum.DONE)
@ -74,4 +91,3 @@ def executor_initializer():
return True

View File

@ -13,19 +13,15 @@ async def on_message(message: AbstractIncomingMessage, queue) -> None:
async def get_messages(inner_queue) -> None:
# Perform connection
async with await connect("amqp://guest:guest@localhost/") as connection:
# Creating a channel
channel = await connection.channel()
await channel.set_qos(prefetch_count=1)
# Declaring queue
queue = await channel.declare_queue(
"hello",
durable=True,
)
# Start listening the queue with name 'task_queue'
await queue.consume(partial(on_message, queue=inner_queue))
print(" [*] Waiting for messages. To exit press CTRL+C")

55
src/core/redis_client.py Normal file
View File

@ -0,0 +1,55 @@
import json
import redis.asyncio as redis
class RedisClient:
SET_NAME = "queue"
TASKS_NAME = "tasks_working"
TASKS_DONE_NAME = "tasks_done"
def __init__(self):
self.connection = redis.Redis(host="localhost", port=6379, db=0)
async def _set_task(self, task: dict) -> int:
async with self.connection as connection:
res = await connection.set(f'{self.TASKS_NAME}:{task["link"]}', json.dumps(task, indent=4).encode('utf-8'))
return res
async def _set_task_done(self, task: dict) -> int:
async with self.connection as connection:
res = await connection.set(
f'{self.TASKS_DONE_NAME}:1:{task["link"]}',
json.dumps(task, indent=4).encode('utf-8')
)
return res
async def _del_task(self, task: dict) -> int:
async with self.connection as connection:
res = await connection.delete(f'{self.TASKS_NAME}:{task["link"]}')
return res
async def set_task_to_queue(self, task: dict) -> int:
async with self.connection as connection:
res = await connection.sadd(self.SET_NAME, json.dumps(task, indent=4).encode('utf-8'))
return res
async def get_queue(self) -> set:
async with self.connection as connection:
res = await connection.smembers(self.SET_NAME)
return res
async def del_task_from_queue(self, task: dict) -> int:
async with self.connection as connection:
res = await connection.srem(self.SET_NAME, json.dumps(task, indent=4).encode('utf-8'))
return res
async def del_task_from_queue_and_add_to_tasks(self, task: dict) -> int:
await self.del_task_from_queue(task)
return await self._set_task(task)
async def del_task_from_tasks_and_add_to_task_done(self, task: dict) -> int:
await self._del_task(task)
return await self._set_task_done(task)

View File

@ -4,10 +4,11 @@ from enum import Enum
class ResultTypeEnum(Enum):
EXCEPTION = "Error"
DONE = "Done"
EXIST = "Exist"
class Result:
def __init__(self, result_type: ResultTypeEnum, value: Exception = None):
def __init__(self, result_type: ResultTypeEnum, value: Exception | bool = None):
self.result_type = result_type
self.value = value

View File

@ -24,12 +24,9 @@ class VideoDownloader:
self.username = username
self.password = password
@staticmethod
def get_unique_video_filename():
prefix = f'vid_{datetime.now().strftime("%Y%m%d%H%M%S")}'
random_uuid4 = uuid.uuid4().hex[:8]
filename = f"{prefix}_{random_uuid4}"
return filename
def get_info(self):
with youtube_dl.YoutubeDL(self.ydl_opts if self.ydl_opts else {}) as ydl:
return ydl.extract_info(self.link, download=False)
def download(self):
domain = urlparse(self.link).netloc