305 lines
11 KiB
Python
305 lines
11 KiB
Python
import asyncio
|
|
from email.utils import formatdate
|
|
import logging
|
|
import os
|
|
import time
|
|
import yt_dlp
|
|
from sqlite3 import Connection
|
|
from src.Config import Config
|
|
from src.Notifier import Notifier
|
|
from src.Download import Download
|
|
from src.DTO.ItemDTO import ItemDTO
|
|
from src.DataStore import DataStore
|
|
from src.Utils import ObjectSerializer, calcDownloadPath, ExtractInfo, mergeConfig
|
|
from datetime import datetime, timezone
|
|
|
|
log = logging.getLogger('DownloadQueue')
|
|
|
|
|
|
class DownloadQueue:
|
|
config: Config = None
|
|
notifier: Notifier = None
|
|
connection: Connection = None
|
|
serializer: ObjectSerializer = None
|
|
queue: DataStore = None
|
|
done: DataStore = None
|
|
event: asyncio.Event = None
|
|
|
|
def __init__(self, config: Config, notifier: Notifier, connection: Connection, serializer: ObjectSerializer):
|
|
self.config = config
|
|
self.notifier = notifier
|
|
self.connection = connection
|
|
self.serializer = serializer
|
|
self.done = DataStore(type='done', config=self.config)
|
|
self.queue = DataStore(type='queue', config=self.config)
|
|
self.done.load()
|
|
self.queue.load()
|
|
|
|
async def initialize(self):
|
|
self.event = asyncio.Event()
|
|
asyncio.create_task(self.__download())
|
|
|
|
async def __add_entry(
|
|
self,
|
|
entry: dict,
|
|
quality: str,
|
|
format: str,
|
|
folder: str,
|
|
ytdlp_config: dict = {},
|
|
ytdlp_cookies: str = '',
|
|
output_template: str = '',
|
|
already=None
|
|
):
|
|
if not entry:
|
|
return {'status': 'error', 'msg': 'Invalid/empty data was given.'}
|
|
|
|
error: str = None
|
|
live_in: str = None
|
|
|
|
if 'live_status' in entry and entry.get('live_status') == 'is_upcoming':
|
|
if 'release_timestamp' in entry and entry.get('release_timestamp'):
|
|
dt_ts = datetime.fromtimestamp(
|
|
entry.get('release_timestamp'), tz=timezone.utc).strftime('%Y-%m-%d %H:%M:%S %z')
|
|
error = f"Live stream is scheduled to start at {dt_ts}"
|
|
live_in = formatdate(entry.get('release_timestamp'))
|
|
else:
|
|
error = 'Live stream not yet started. And no date is set.'
|
|
else:
|
|
error = entry['msg'] if 'msg' in entry else None
|
|
|
|
etype = entry.get('_type') or 'video'
|
|
if etype == 'playlist':
|
|
entries = entry['entries']
|
|
log.info(f'playlist detected with {len(entries)} entries')
|
|
playlist_index_digits = len(str(len(entries)))
|
|
results = []
|
|
for index, etr in enumerate(entries, start=1):
|
|
etr['playlist'] = entry['id']
|
|
etr['playlist_index'] = '{{0:0{0:d}d}}'.format(
|
|
playlist_index_digits).format(index)
|
|
for property in ('id', 'title', 'uploader', 'uploader_id'):
|
|
if property in entry:
|
|
etr[f'playlist_{property}'] = entry[property]
|
|
|
|
results.append(
|
|
await self.__add_entry(
|
|
entry=etr,
|
|
quality=quality,
|
|
format=format,
|
|
folder=folder,
|
|
ytdlp_config=ytdlp_config,
|
|
ytdlp_cookies=ytdlp_cookies,
|
|
output_template=output_template,
|
|
already=already
|
|
)
|
|
)
|
|
|
|
if any(res['status'] == 'error' for res in results):
|
|
return {'status': 'error', 'msg': ', '.join(res['msg'] for res in results if res['status'] == 'error' and 'msg' in res)}
|
|
|
|
return {'status': 'ok'}
|
|
elif (etype == 'video' or etype.startswith('url')) and 'id' in entry and 'title' in entry:
|
|
|
|
logging.debug(
|
|
f"entry: {entry.get('id', None)} - {entry.get('webpage_url', None)} - {entry.get('url', None) }")
|
|
|
|
if self.done.exists(key=entry['id'], url=entry.get('webpage_url') or entry['url']):
|
|
item = self.done.get(key=entry['id'], url=entry.get(
|
|
'webpage_url') or entry['url'])
|
|
|
|
logging.debug(
|
|
f'Item [{item.info.title}] already downloaded. Removing from history.')
|
|
|
|
await self.clear([item.info._id])
|
|
|
|
if self.queue.exists(key=entry['id'], url=entry.get('webpage_url') or entry['url']):
|
|
logging.info(
|
|
f'Item [{item.info.title}] already in download queue')
|
|
return {'status': 'error', 'msg': 'Link already queued for downloading.'}
|
|
|
|
dl = ItemDTO(
|
|
id=entry['id'],
|
|
title=entry['title'],
|
|
url=entry.get('webpage_url') or entry['url'],
|
|
quality=quality,
|
|
format=format,
|
|
folder=folder,
|
|
ytdlp_cookies=ytdlp_cookies,
|
|
ytdlp_config=ytdlp_config,
|
|
output_template=output_template if output_template else self.config.output_template,
|
|
datetime=formatdate(time.time()),
|
|
error=error,
|
|
is_live=entry['is_live'] if 'is_live' in entry else None,
|
|
live_in=live_in,
|
|
)
|
|
|
|
try:
|
|
dldirectory = calcDownloadPath(
|
|
basePath=self.config.download_path,
|
|
folder=folder
|
|
)
|
|
except Exception as e:
|
|
return {'status': 'error', 'msg': str(e)}
|
|
|
|
output_chapter = self.config.output_template_chapter
|
|
|
|
for property, value in entry.items():
|
|
if property.startswith("playlist"):
|
|
dl.output_template = dl.output_template.replace(
|
|
f"%({property})s", str(value))
|
|
|
|
itemDownload = self.queue.put(Download(
|
|
info=dl,
|
|
download_dir=dldirectory,
|
|
temp_dir=self.config.temp_path,
|
|
output_template_chapter=output_chapter,
|
|
default_ytdl_opts=self.config.ytdl_options,
|
|
debug=bool(self.config.ytdl_debug)
|
|
))
|
|
|
|
self.event.set()
|
|
await self.notifier.added(itemDownload.info)
|
|
|
|
return {
|
|
'status': 'ok'
|
|
}
|
|
elif etype.startswith('url'):
|
|
return await self.add(
|
|
entry=entry['url'],
|
|
quality=quality,
|
|
format=format,
|
|
folder=folder,
|
|
ytdlp_config=ytdlp_config,
|
|
ytdlp_cookies=ytdlp_cookies,
|
|
output_template=output_template,
|
|
already=already
|
|
)
|
|
|
|
return {
|
|
'status': 'error',
|
|
'msg': f'Unsupported resource "{etype}"'
|
|
}
|
|
|
|
async def add(
|
|
self,
|
|
url: str,
|
|
quality: str,
|
|
format: str,
|
|
folder: str,
|
|
ytdlp_config: dict = {},
|
|
ytdlp_cookies: str = '',
|
|
output_template: str = '',
|
|
already=None
|
|
):
|
|
ytdlp_config = ytdlp_config if ytdlp_config else {}
|
|
|
|
log.info(
|
|
f'adding {url}: {quality=} {format=} {folder=} {output_template=} {ytdlp_cookies=} {ytdlp_config=}')
|
|
|
|
already = set() if already is None else already
|
|
if url in already:
|
|
log.info('recursion detected, skipping')
|
|
return {'status': 'ok'}
|
|
else:
|
|
already.add(url)
|
|
try:
|
|
entry = await asyncio.get_running_loop().run_in_executor(
|
|
None,
|
|
ExtractInfo,
|
|
mergeConfig(self.config.ytdl_options, ytdlp_config),
|
|
url,
|
|
bool(self.config.ytdl_debug)
|
|
)
|
|
if not entry:
|
|
return {
|
|
'status': 'error',
|
|
'msg': 'No metadata, most likely video has been downloaded before.' if self.config.keep_archive else 'Unable to extract info check logs.'
|
|
}
|
|
logging.debug(f'entry: extract info says: {entry}')
|
|
except yt_dlp.utils.YoutubeDLError as exc:
|
|
return {'status': 'error', 'msg': str(exc)}
|
|
|
|
return await self.__add_entry(
|
|
entry=entry,
|
|
quality=quality,
|
|
format=format,
|
|
folder=folder,
|
|
ytdlp_config=ytdlp_config,
|
|
ytdlp_cookies=ytdlp_cookies,
|
|
output_template=output_template,
|
|
already=already,
|
|
)
|
|
|
|
async def cancel(self, ids):
|
|
for id in ids:
|
|
if not self.queue.exists(id):
|
|
log.warn(f'requested cancel for non-existent download {id}')
|
|
continue
|
|
if self.queue.get(id).started():
|
|
self.queue.get(id).cancel()
|
|
else:
|
|
self.queue.delete(id)
|
|
await self.notifier.canceled(id)
|
|
|
|
return {'status': 'ok'}
|
|
|
|
async def clear(self, ids):
|
|
for id in ids:
|
|
if not self.done.exists(key=id):
|
|
log.warn(f'requested delete for non-existent download {id}')
|
|
continue
|
|
self.done.delete(id)
|
|
await self.notifier.cleared(id)
|
|
|
|
return {'status': 'ok'}
|
|
|
|
def get(self) -> dict[str, list[dict[str, ItemDTO]]]:
|
|
items = {'queue': {}, 'done': {}}
|
|
|
|
for k, v in self.queue.saved_items():
|
|
items['queue'][k] = v
|
|
|
|
for k, v in self.done.saved_items():
|
|
items['done'][k] = v
|
|
|
|
for k, v in self.queue.items():
|
|
if not k in items['queue']:
|
|
items['queue'][k] = v.info
|
|
|
|
for k, v in self.done.items():
|
|
if not k in items['done']:
|
|
items['done'][k] = v.info
|
|
|
|
return items
|
|
|
|
async def __download(self):
|
|
while True:
|
|
while self.queue.empty():
|
|
log.info('waiting for item to download.')
|
|
await self.event.wait()
|
|
self.event.clear()
|
|
|
|
id, entry = self.queue.next()
|
|
log.info(
|
|
f'downloading {id=} - {entry.info.title=} - {entry.info.url=} - {entry.info.folder=}')
|
|
|
|
await entry.start(self.notifier)
|
|
|
|
if entry.info.status != 'finished':
|
|
if entry.tmpfilename and os.path.isfile(entry.tmpfilename):
|
|
try:
|
|
os.remove(entry.tmpfilename)
|
|
except:
|
|
pass
|
|
|
|
entry.info.status = 'error'
|
|
|
|
entry.close()
|
|
|
|
if self.queue.exists(id):
|
|
self.queue.delete(id)
|
|
if entry.canceled:
|
|
await self.notifier.canceled(id)
|
|
else:
|
|
self.done.put(entry)
|
|
await self.notifier.completed(entry.info)
|