457 lines
14 KiB
Python
457 lines
14 KiB
Python
import asyncio
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import pkgutil
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from aiohttp import web
|
|
|
|
from .config import Config
|
|
from .DownloadQueue import DownloadQueue
|
|
from .encoder import Encoder
|
|
from .Events import EventBus, Events
|
|
from .ItemDTO import Item
|
|
from .Scheduler import Scheduler
|
|
from .Services import Services
|
|
from .Singleton import Singleton
|
|
from .Utils import init_class, validate_url
|
|
|
|
LOG: logging.Logger = logging.getLogger("tasks")
|
|
|
|
|
|
@dataclass(kw_only=True)
|
|
class Task:
|
|
id: str
|
|
name: str
|
|
url: str
|
|
folder: str = ""
|
|
preset: str = ""
|
|
timer: str = ""
|
|
template: str = ""
|
|
cli: str = ""
|
|
auto_start: bool = True
|
|
|
|
def serialize(self) -> dict:
|
|
return self.__dict__
|
|
|
|
def json(self) -> str:
|
|
return Encoder().encode(self.serialize())
|
|
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
return self.serialize().get(key, default)
|
|
|
|
|
|
class Tasks(metaclass=Singleton):
|
|
"""
|
|
This class is used to manage the tasks.
|
|
"""
|
|
|
|
_tasks: list[Task] = []
|
|
"""The tasks."""
|
|
|
|
_instance = None
|
|
"""The instance of the Tasks."""
|
|
|
|
def __init__(
|
|
self,
|
|
file: str | None = None,
|
|
loop: asyncio.AbstractEventLoop | None = None,
|
|
config: Config | None = None,
|
|
encoder: Encoder | None = None,
|
|
scheduler: Scheduler | None = None,
|
|
):
|
|
Tasks._instance = self
|
|
|
|
config = config or Config.get_instance()
|
|
|
|
self._debug: bool = config.debug
|
|
self._default_preset: str = config.default_preset
|
|
self._file: Path = Path(file) if file else Path(config.config_path).joinpath("tasks.json")
|
|
self._encoder: Encoder = encoder or Encoder()
|
|
self._loop: asyncio.AbstractEventLoop = loop or asyncio.get_event_loop()
|
|
self._scheduler: Scheduler = scheduler or Scheduler.get_instance()
|
|
self._notify: EventBus = EventBus.get_instance()
|
|
self._task_handler = HandleTask(self._scheduler, self, config)
|
|
self._downloadQueue = DownloadQueue.get_instance()
|
|
|
|
if self._file.exists() and "600" != self._file.stat().st_mode:
|
|
try:
|
|
self._file.chmod(0o600)
|
|
except Exception:
|
|
pass
|
|
|
|
@staticmethod
|
|
def get_instance() -> "Tasks":
|
|
"""
|
|
Get the instance of the class.
|
|
|
|
Returns:
|
|
Tasks: The instance of the class.
|
|
|
|
"""
|
|
if not Tasks._instance:
|
|
Tasks._instance = Tasks()
|
|
|
|
return Tasks._instance
|
|
|
|
async def on_shutdown(self, _: web.Application):
|
|
pass
|
|
|
|
def attach(self, _: web.Application):
|
|
"""
|
|
Attach the tasks to the aiohttp application.
|
|
|
|
Args:
|
|
_ (web.Application): The aiohttp application.
|
|
|
|
"""
|
|
self.load()
|
|
self._notify.subscribe(
|
|
Events.TASKS_ADD,
|
|
lambda data, _, **kwargs: self.save(**data.data), # noqa: ARG005
|
|
f"{__class__.__name__}.add",
|
|
)
|
|
|
|
def get_all(self) -> list[Task]:
|
|
"""Return the tasks."""
|
|
return self._tasks
|
|
|
|
def load(self) -> "Tasks":
|
|
"""
|
|
Load the tasks.
|
|
|
|
Returns:
|
|
Tasks: The current instance.
|
|
|
|
"""
|
|
has_tasks: bool = len(self._tasks) > 0
|
|
self.clear()
|
|
|
|
if not self._file.exists() or self._file.stat().st_size < 1:
|
|
return self
|
|
|
|
try:
|
|
LOG.info(f"{'Reloading' if has_tasks else 'Loading'} '{self._file}'.")
|
|
tasks = json.loads(self._file.read_text())
|
|
except Exception as e:
|
|
LOG.error(f"Error loading '{self._file}'. '{e!s}'.")
|
|
return self
|
|
|
|
if not tasks or len(tasks) < 1:
|
|
return self
|
|
|
|
for i, task in enumerate(tasks):
|
|
try:
|
|
Tasks.validate(task)
|
|
task: Task = init_class(Task, task)
|
|
except Exception as e:
|
|
LOG.error(f"Failed to parse task at list position '{i}'. '{e!s}'.")
|
|
continue
|
|
|
|
self._tasks.append(task)
|
|
|
|
if not task.timer or "[only_handler]" in task.name:
|
|
continue
|
|
|
|
try:
|
|
self._scheduler.add(timer=task.timer, func=self._runner, args=(task,), id=task.id)
|
|
|
|
try:
|
|
from cronsim import CronSim
|
|
|
|
cs = CronSim(task.timer, datetime.now(UTC))
|
|
schedule_time: str = cs.explain()
|
|
except Exception:
|
|
schedule_time = task.timer
|
|
|
|
if not has_tasks:
|
|
LOG.info(f"Task '{task.name}' queued to be executed '{schedule_time}'.")
|
|
except Exception as e:
|
|
LOG.exception(e)
|
|
LOG.error(f"Failed to queue '{i}: {task.name}'. '{e!s}'.")
|
|
|
|
return self
|
|
|
|
def clear(self, shutdown: bool = False) -> "Tasks":
|
|
"""
|
|
Clear all tasks.
|
|
|
|
Returns:
|
|
Tasks: The current instance.
|
|
|
|
"""
|
|
if len(self._tasks) < 1:
|
|
return self
|
|
|
|
for task in self._tasks:
|
|
if not self._scheduler.has(task.id):
|
|
continue
|
|
|
|
try:
|
|
LOG.debug(f"Stopping '{task.name}'.")
|
|
self._scheduler.remove(task.id)
|
|
except Exception as e:
|
|
if not shutdown:
|
|
LOG.exception(e)
|
|
LOG.error(f"Failed to stop '{task.name}'. '{e!s}'.")
|
|
|
|
self._tasks.clear()
|
|
|
|
return self
|
|
|
|
@staticmethod
|
|
def validate(task: Task | dict) -> bool:
|
|
"""
|
|
Validate the task.
|
|
|
|
Args:
|
|
task (Task|dict): The task to validate.
|
|
|
|
Returns:
|
|
bool: True if the task is valid, False otherwise.
|
|
|
|
"""
|
|
if not isinstance(task, dict):
|
|
if not isinstance(task, Task):
|
|
msg = "Invalid task type."
|
|
raise ValueError(msg) # noqa: TRY004
|
|
|
|
task = task.serialize()
|
|
|
|
if not task.get("name"):
|
|
msg = "No name found."
|
|
raise ValueError(msg)
|
|
|
|
task["name"] = task["name"].strip()
|
|
|
|
if not task.get("url"):
|
|
msg = "No URL found."
|
|
raise ValueError(msg)
|
|
|
|
task["url"] = task["url"].strip()
|
|
try:
|
|
validate_url(task["url"], allow_internal=True)
|
|
except ValueError as e:
|
|
msg = f"Invalid URL format. '{e!s}'."
|
|
raise ValueError(msg) from e
|
|
|
|
if task.get("timer"):
|
|
try:
|
|
from cronsim import CronSim
|
|
|
|
CronSim(task.get("timer"), datetime.now(UTC))
|
|
task["timer"] = str(task["timer"]).strip()
|
|
except Exception as e:
|
|
msg = f"Invalid timer format. '{e!s}'."
|
|
raise ValueError(msg) from e
|
|
|
|
if task.get("cli"):
|
|
try:
|
|
from .Utils import arg_converter
|
|
|
|
arg_converter(args=task.get("cli"))
|
|
task["cli"] = str(task["cli"]).strip()
|
|
except Exception as e:
|
|
msg = f"Invalid command options for yt-dlp. '{e!s}'."
|
|
raise ValueError(msg) from e
|
|
|
|
return True
|
|
|
|
def save(self, tasks: list[Task | dict]) -> "Tasks":
|
|
"""
|
|
Save the tasks.
|
|
|
|
Args:
|
|
tasks (list[Task]): The tasks to save.
|
|
|
|
Returns:
|
|
Tasks: The current instance.
|
|
|
|
"""
|
|
for i, task in enumerate(tasks):
|
|
try:
|
|
self.validate(task)
|
|
|
|
if not isinstance(task, Task):
|
|
task: Task = init_class(Task, task)
|
|
tasks[i] = task
|
|
except ValueError as e:
|
|
LOG.error(f"Failed to validate item '{i}: {task.name}'. '{e}'.")
|
|
continue
|
|
except Exception as e:
|
|
LOG.error(f"Failed to save task '{i}'. '{e!s}'.")
|
|
continue
|
|
|
|
try:
|
|
self._file.write_text(json.dumps([i.serialize() for i in tasks], indent=4))
|
|
LOG.info(f"Updated '{self._file}'.")
|
|
except Exception as e:
|
|
LOG.error(f"Error saving '{self._file}'. '{e!s}'.")
|
|
|
|
return self
|
|
|
|
async def _runner(self, task: Task) -> None:
|
|
"""
|
|
Run the task.
|
|
|
|
Args:
|
|
task (Task): The task to run.
|
|
|
|
Returns:
|
|
None
|
|
|
|
"""
|
|
timeNow: str = datetime.now(UTC).isoformat()
|
|
try:
|
|
started: float = time.time()
|
|
if not task.url:
|
|
LOG.error(f"Failed to dispatch '{task.name}'. No URL found.")
|
|
return
|
|
|
|
preset: str = str(task.preset or self._default_preset)
|
|
folder: str = task.folder if task.folder else ""
|
|
template: str = task.template if task.template else ""
|
|
cli: str = task.cli if task.cli else ""
|
|
|
|
status = await self._downloadQueue.add(
|
|
item=Item.format(
|
|
{
|
|
"url": task.url,
|
|
"preset": preset,
|
|
"folder": folder,
|
|
"template": template,
|
|
"cli": cli,
|
|
}
|
|
)
|
|
)
|
|
|
|
timeNow = datetime.now(UTC).isoformat()
|
|
ended: float = time.time()
|
|
LOG.info(f"Task '{task.name}' completed at '{timeNow}' took '{ended - started:.2f}' seconds.")
|
|
|
|
_tasks = [
|
|
self._notify.emit(
|
|
Events.TASK_DISPATCHED,
|
|
data=status,
|
|
title=f"Task '{task.name}' dispatched",
|
|
message=f"Task '{task.name}' dispatched at '{timeNow}'.",
|
|
),
|
|
self._notify.emit(
|
|
Events.LOG_SUCCESS,
|
|
data={"lowPriority": True},
|
|
title="Task completed",
|
|
message=f"Task '{task.name}' completed in '{ended - started:.2f}'.",
|
|
),
|
|
]
|
|
|
|
await asyncio.gather(*_tasks)
|
|
except Exception as e:
|
|
LOG.error(f"Failed to execute '{task.name}' at '{timeNow}'. '{e!s}'.")
|
|
await self._notify.emit(
|
|
Events.LOG_ERROR,
|
|
title="Task failed",
|
|
message=f"Failed to execute '{task.name}'. '{e!s}'",
|
|
)
|
|
|
|
|
|
class HandleTask:
|
|
_tasks: Tasks
|
|
|
|
def __init__(self, scheduler: Scheduler, tasks: Tasks, config: Config) -> None:
|
|
self._tasks = tasks
|
|
self._handlers: list[type] = self._discover()
|
|
|
|
timer = config.tasks_handler_timer
|
|
try:
|
|
from cronsim import CronSim
|
|
|
|
CronSim(timer, datetime.now(UTC))
|
|
except Exception as e:
|
|
timer = "15 */1 * * *"
|
|
LOG.error(f"Invalid timer format. '{e!s}'. Defaulting to '{timer}'.")
|
|
|
|
scheduler.add(
|
|
timer=timer,
|
|
func=self._dispatcher,
|
|
id=f"{__class__.__name__}._dispatcher",
|
|
)
|
|
|
|
def _dispatcher(self):
|
|
for task in self._tasks.get_all():
|
|
if "[no_handler]" in task.name:
|
|
continue
|
|
|
|
if not task.timer and "[only_handler]" not in task.name:
|
|
continue
|
|
|
|
try:
|
|
handler = self._find_handler(task)
|
|
if handler is None:
|
|
continue
|
|
|
|
coro = self.dispatch(task, handler=handler)
|
|
t = asyncio.create_task(coro, name=f"task-{task.id}")
|
|
t.add_done_callback(lambda fut, t=task: self._handle_exception(fut, t))
|
|
except Exception as e:
|
|
LOG.error(f"Failed to handle task '{task.name}'. '{e!s}'.")
|
|
|
|
def _handle_exception(self, fut: asyncio.Task, task: Task) -> None:
|
|
if fut.cancelled():
|
|
return
|
|
|
|
if exc := fut.exception():
|
|
LOG.error(f"Exception while handling task '{task.name}': {exc}")
|
|
|
|
def _find_handler(self, task: Task) -> type | None:
|
|
for cls in self._handlers:
|
|
try:
|
|
if Services.get_instance().handle_sync(handler=cls.can_handle, task=task):
|
|
return cls
|
|
except Exception as e:
|
|
LOG.exception(e)
|
|
continue
|
|
|
|
return None
|
|
|
|
async def dispatch(self, task: Task, handler: type | None = None, **kwargs) -> Any | None:
|
|
"""
|
|
Dispatch a task to the appropriate handler.
|
|
|
|
Args:
|
|
task (Task): The task to dispatch.
|
|
handler (type|None): Optional specific handler to use instead of finding one.
|
|
**kwargs: Additional context to pass to the handler.
|
|
|
|
Returns:
|
|
Any|None: The result of the handler's execution, or None if no handler found.
|
|
|
|
"""
|
|
if not handler:
|
|
handler = self._find_handler(task)
|
|
if handler is None:
|
|
return None
|
|
|
|
return await Services.get_instance().handle_async(handler=handler.handle, task=task, **kwargs)
|
|
|
|
def _discover(self) -> list[type]:
|
|
import importlib
|
|
|
|
import app.library.task_handlers as handlers_pkg
|
|
|
|
handlers: list[type] = []
|
|
|
|
for _, module_name, _ in pkgutil.iter_modules(handlers_pkg.__path__):
|
|
module = importlib.import_module(f"{handlers_pkg.__name__}.{module_name}")
|
|
for _, cls in inspect.getmembers(module, inspect.isclass):
|
|
if cls.__module__ != module.__name__:
|
|
continue
|
|
|
|
if callable(getattr(cls, "can_handle", None)) and callable(getattr(cls, "handle", None)):
|
|
handlers.append(cls)
|
|
|
|
return handlers
|