autotako

job_render.py (19536B)

---

 #!/usr/bin/python3
 
 import asyncio
 import datetime
 import enum
 import itertools
 import logging
 import pathlib
 import re
 from typing import Self
 
 import gofile.api  # type: ignore
 import httpx
 import microdot  # type: ignore
 import msgspec
 import qbittorrentapi  # type: ignore
 import stamina
 import torf  # type: ignore
 
 from .config import WebDavConfig, config_ctx
 from .database import database_ctx
 
 app = microdot.Microdot()
 
 log = logging.getLogger(__name__)
 
 background_tasks = set()
 render_tasks: dict[str, asyncio.Task] = {}
 upload_tasks: dict[pathlib.Path, asyncio.Task] = {}
 
 # https://www.japanesewithanime.com/2017/05/quotation-marks-japanese.html#sumitsukikakko
 # strip extraneous spaces and hashtag
 _pattern_bracketed_kw = re.compile(r"^【[\s#]*(?P<topic>.*?)[\s]?】")
 
 # bae prefers the double angle brackets because she's quirky like that
 _pattern_angled_bracketed_kw = re.compile(r"^≪[\s#]*(?P<topic>.*?)[\s]?≫")
 
 
 class JobUploadCriteria(enum.StrEnum):
     MANUALLY = "manual"
     WHEN_FINISHED = "finished"
     IF_PRIVATE = "private"
     IF_CUT = "cut"
 
     @property
     def automatic(self) -> bool:
         return self != JobUploadCriteria.MANUALLY
 
 
 class MoomboxDownloadStatus(enum.StrEnum):
     UNKNOWN = "Unknown"
     UNAVAILABLE = "Unavailable"
     WAITING = "Waiting"
     DOWNLOADING = "Downloading"
     MUXING = "Muxing"
     FINISHED = "Finished"
     ERROR = "Error"
 
 
 class MoomboxJobInfo(msgspec.Struct, kw_only=True, frozen=True):
     id: str
     title: str | None
     author: str | None
     channel_id: str | None
     video_id: str | None
     manifest_progress: dict
     status: MoomboxDownloadStatus = MoomboxDownloadStatus.UNKNOWN
     output_paths: set[str] | None = msgspec.field(default_factory=set)
     scheduled_start_datetime: datetime.datetime | None = None
 
     @property
     def uploadability_state(self) -> JobUploadCriteria | None:
         # returns a JobUploadCriteria that can be tested for equality to run
         if self.status != MoomboxDownloadStatus.FINISHED:
             return None
         if not self.files_available:
             return None
         if not self.is_configured_for_upload:
             return None
         return JobUploadCriteria.WHEN_FINISHED
 
     @property
     def is_configured_for_upload(self) -> bool:
         if not self.channel_id:
             return False
         config = config_ctx.get()
         channel_config = config.get_channel_config_by_id(self.channel_id)
         if not channel_config:
             return False
         return True
 
     @property
     def files_available(self) -> bool:
         return bool(self.output_pathobjs) and all(p.exists() for p in self.output_pathobjs)
 
     @property
     def output_pathobjs(self) -> set[pathlib.Path]:
         return set(map(pathlib.Path, self.output_paths or []))
 
 
 def _extract_name_from_stream_title(title: str) -> str | None:
     """
     Attempts to extract a preset name from the stream title.
     """
     match = _pattern_bracketed_kw.search(title) or _pattern_angled_bracketed_kw.search(title)
     if match:
         return str(match["topic"]).strip()
     return None
 
 
 class JobConfig(msgspec.Struct):
     name: str
     upload: JobUploadCriteria
     author_note: str | None = None
 
     @classmethod
     def from_moombox_job(cls, job: MoomboxJobInfo) -> Self:
         # generate a default config from the job
         name = "Unspecified Stream"
         upload = JobUploadCriteria.MANUALLY
 
         if not job.title:
             return cls(name, upload)
 
         name = _extract_name_from_stream_title(job.title) or name
 
         # mapping of phrases and any possible synonyms
         # if a synonym isn't present then an empty set should be present
         phrases = {
             "unarchived": {"no archive", "unarchive"},
             "karaoke": {"rock n' rawr", "歌枠", "sing"},
             "rebroadcast": {"kara-rewind"},
         }
 
         phrase_set = set(phrases.keys()) | set(itertools.chain.from_iterable(phrases.values()))
         # not sure why mypy complains about job.title being possibly None here if left bare...
         matched_phrases = set(
             filter(lambda term: term in (job.title.lower() if job.title else ""), phrase_set)
         )
 
         # dedupe synonymous terms
         for phrase, syn_phrases in phrases.items():
             if syn_phrases and matched_phrases & syn_phrases:
                 matched_phrases = (matched_phrases - syn_phrases) | {phrase}
 
         if matched_phrases == {"unarchived"}:
             name = "Unarchived Stream"
             upload = JobUploadCriteria.WHEN_FINISHED
         elif matched_phrases == {"karaoke"}:
             name = "Karaoke"
             upload = JobUploadCriteria.IF_CUT
         elif matched_phrases == {"unarchived", "karaoke"}:
             name = "Unarchived Karaoke"
             upload = JobUploadCriteria.WHEN_FINISHED
         elif matched_phrases == {"unarchived", "karaoke", "rebroadcast"}:
             name = "Unarchived Karaoke (rebroadcast)"
             upload = JobUploadCriteria.MANUALLY
 
         # cannot auto-upload if no channel configuration exists
         if not job.is_configured_for_upload:
             upload = JobUploadCriteria.MANUALLY
 
         return cls(name, upload)
 
 
 def stream_display_date(dt: datetime.datetime) -> datetime.date:
     """
     Given the stream start time, returns a display date.
 
     If the stream start time is within the last 10 minutes of the day, we treat the next day
     as the stream date instead based on the assumption that the streamer has a waiting screen
     set up.
     """
     end_of_day = (23, 50)
     return dt.date() + datetime.timedelta(days=1 if (dt.hour, dt.minute) >= end_of_day else 0)
 
 
 async def get_moombox_job_by_id(jobid: str) -> MoomboxJobInfo | None:
     config = config_ctx.get()
     async with httpx.AsyncClient() as client:
         for attempt in stamina.retry_context(
             on=httpx.HTTPError, attempts=None, timeout=None, wait_initial=0.5, wait_max=10.0
         ):
             with attempt:
                 result = await client.get(f"{config.moombox_url}/status")
         for job in result.json():
             if job["id"] == jobid:
                 return msgspec.convert(job, type=MoomboxJobInfo)
         return None
 
 
 async def gofile_upload_and_persist(
     cursor, upload_file: pathlib.Path, job_id: str, token: str, folder_id: str
 ):
     with upload_file.open("rb") as upload_filereader:
         await gofile.api.upload_single(upload_filereader, token, folder_id)
         database = database_ctx.get()
         cursor = database.cursor()
         cursor.execute(
             "INSERT INTO gofile_files VALUES (?, ?)",
             (job_id, str(upload_file)),
         )
         database.commit()
 
 
 async def do_gofile_upload(job: MoomboxJobInfo):
     """
     Prepares a guest upload for the specific job's output files, then returns the folder URL
     while the uploads are running in the background.
 
     This operation is idempotent; if it is interrupted, running this function will only upload
     files that weren't already.
     """
     database = database_ctx.get()
 
     known_files = set()
     for known_file, *_ in database.execute(
         "SELECT filepath FROM gofile_files WHERE job_id = ?", (job.id,)
     ):
         known_files.add(pathlib.Path(known_file))
 
     pending_files = set(map(pathlib.Path, job.output_paths or [])) - known_files
 
     token = None
     folder_id = None
     gofile_url = None
 
     cursor = database.cursor()
     cursor.execute(
         "SELECT token, folder_id, url FROM gofile_tokens WHERE job_id = ?", (job.id,)
     )
     job_token_info = cursor.fetchone()
     if job_token_info:
         token, folder_id, gofile_url = job_token_info
     else:
         # get the smallest file to quickly get the upload result
         pending_file = min(pending_files, key=lambda p: p.stat().st_size)
         if pending_file in upload_tasks:
             # another task is uploading the same file; just block until that one is available
             upload = await upload_tasks[pending_file]
             token, folder_id, gofile_url = (
                 upload.result.guest_token,
                 upload.result.parent_folder,
                 upload.result.download_page,
             )
         else:
             with pending_file.open("rb") as fo:
                 upload_tasks[pending_file] = gofile.api.upload_single(fo)
                 upload = await upload_tasks[pending_file]
                 token, folder_id, gofile_url = (
                     upload.result.guest_token,
                     upload.result.parent_folder,
                     upload.result.download_page,
                 )
             cursor.execute(
                 "INSERT INTO gofile_tokens VALUES (?, ?, ?, ?)",
                 (job.id, token, folder_id, gofile_url),
             )
             cursor.execute(
                 "INSERT INTO gofile_files VALUES (?, ?)",
                 (job.id, str(pending_file)),
             )
             database.commit()
         pending_files.discard(pending_file)
 
     # upload any remaining files in the background
     for pending_file in pending_files:
         # don't double-dip and upload an in-progress file
         if pending_file in upload_tasks:
             continue
         upload_tasks[pending_file] = asyncio.create_task(
             gofile_upload_and_persist(cursor, pending_file, job.id, token, folder_id)
         )
     return gofile_url
 
 
 async def do_webdav_upload(webdav: WebDavConfig, filepath: pathlib.Path, target: str):
     auth = httpx.BasicAuth(username=webdav.username, password=webdav.password)
     async with httpx.AsyncClient(auth=auth) as client:
         connection_warning_seen = False
         while True:
             try:
                 dest = f"{webdav.base_url}/{target}"
                 file_check = await client.head(dest)
                 if file_check.status_code == httpx.codes.NOT_FOUND:
                     with filepath.open("rb") as fh:
                         await client.put(dest, content=fh.read())
                 return
             except httpx.ConnectTimeout:
                 if not connection_warning_seen:
                     log.warning(f"Failed to connect to {webdav.base_url}.  Retrying...")
                     connection_warning_seen = True
                 await asyncio.sleep(10)
 
 
 async def _process_job(jobid):
     job = await get_moombox_job_by_id(jobid)
     if not job:
         return "Couldn't find matching job", 404
 
     # TODO fill in necessary fields in template, etc.
     job_config = _get_saved_jobconf(job.id) or JobConfig.from_moombox_job(job)
 
     config = config_ctx.get()
 
     channel = config.get_channel_config_by_id(job.channel_id)
     stream_time = job.scheduled_start_datetime or datetime.datetime.now(tz=datetime.UTC)
 
     torrent_output_dir = pathlib.Path("torrents")
     torrent_output_dir.mkdir(exist_ok=True)
 
     display_date = stream_display_date(stream_time)
     display_date_str = display_date.strftime("%Y%m%d")
     torrent_name = f"[{display_date_str}] {job_config.name} ({job.video_id})"
 
     folder_name = None
     if stream_time:
         folder_name = stream_time.strftime("%Y-%m-%dT%H:%MZ")
 
     readme_finalized = False
     render_kwargs = {}
     if job.output_paths:
         torrent_file = torrent_output_dir / f"{job.id} ({job.video_id}).torrent"
         source_files = list(map(pathlib.Path, job.output_paths))
 
         if not all(f.exists() for f in source_files):
             return
 
         save_path = str(source_files[0].parent)
 
         if not torrent_file.exists():
             t = torf.Torrent(
                 trackers=config.torrent.trackers,
             )
 
             # TODO extract key phrases from title / description
 
             t.filepaths = source_files
             # note that the name also sets the subdirectory for the files
             # as such, using other names for the same files *will* change the infohash
             t.name = torrent_name
             success = await asyncio.to_thread(t.generate)
             if not success:
                 raise RuntimeError("Failed to generate torrent")
             t.write(torrent_file)
             t.write(torrent_output_dir / f"{torrent_name}.torrent")
 
             # send torrent to qbittorrent for seeding
             if config.qbittorrent:
                 if config.qbittorrent.default_save_path:
                     save_path = config.qbittorrent.default_save_path
                 qbt_client = qbittorrentapi.Client(
                     host=config.qbittorrent.host,
                     username=config.qbittorrent.username,
                     password=config.qbittorrent.password,
                     VERIFY_WEBUI_CERTIFICATE=False,
                     RAISE_NOTIMPLEMENTEDERROR_FOR_UNIMPLEMENTED_API_ENDPOINTS=True,
                 )
                 task = asyncio.create_task(
                     asyncio.to_thread(
                         qbt_client.torrents_add,
                         torrent_files=torrent_file,
                         save_path=save_path,
                         content_layout="NoSubfolder",
                         is_paused=config.qbittorrent.start_paused,
                     )
                 )
                 background_tasks.add(task)
                 task.add_done_callback(background_tasks.discard)
         else:
             t = torf.Torrent.read(torrent_file)
         log.debug(t)
         log.debug(t.magnet(size=False))
         log.debug(t.files)
 
         # punt file to webdav remote
         if config.webdav and channel and channel.webdav_path:
             target_base = f"{channel.webdav_path}/{folder_name}"
             task = asyncio.create_task(
                 do_webdav_upload(
                     config.webdav, torrent_file, f"{target_base}/{torrent_name}.torrent"
                 )
             )
             background_tasks.add(task)
             task.add_done_callback(background_tasks.discard)
 
         render_kwargs["magnet_url"] = t.magnet(size=False)
 
         # punt files to gofile
         render_kwargs["gofile_url"] = await do_gofile_upload(job)
         readme_finalized = True
     rendered_job = await microdot.jinja.Template("job.md").render_async(
         job=job,
         author_override=channel.name if channel else None,
         stream_time=stream_time.strftime("%Y-%m-%dT%H:%M:%SZ") if stream_time else "(unknown)",
         config=job_config,
         **render_kwargs,
     )
     if readme_finalized and config.webdav and channel and channel.webdav_path:
         rendered_job_file = torrent_output_dir / f"{job.id} ({job.video_id}).md"
         rendered_job_file.write_text(rendered_job)
         target_base = f"{channel.webdav_path}/{folder_name}"
         task = asyncio.create_task(
             do_webdav_upload(config.webdav, rendered_job_file, f"{target_base}/REPORT.md")
         )
         background_tasks.add(task)
         task.add_done_callback(background_tasks.discard)
 
         database = database_ctx.get()
         cursor = database.cursor()
         cursor.execute(
             "INSERT INTO job_status VALUES (?, ?) ON CONFLICT(job_id) DO UPDATE SET completed=excluded.completed",
             (jobid, True),
         )
         database.commit()
     if not readme_finalized and jobid in render_tasks:
         # allow rerun if not finalized
         del render_tasks[jobid]
     return rendered_job
 
 
 def get_process_job_task(jobid: str) -> asyncio.Task:
     """
     Creates or retrieves a singleton job rendering task.
     This is decoupled from the request callback to allow cancellations and multiple consumers.
     """
     if jobid in render_tasks and render_tasks[jobid].done() and render_tasks[jobid].exception():
         del render_tasks[jobid]
     if jobid not in render_tasks:
         render_tasks[jobid] = asyncio.create_task(_process_job(jobid))
     return render_tasks[jobid]
 
 
 def _get_saved_jobconf(jobid: str) -> JobConfig | None:
     database = database_ctx.get()
     cursor = database.cursor()
     cursor.execute("SELECT config FROM job_config WHERE job_id = ?", (jobid,))
     result = cursor.fetchone()
     if result:
         job_config_payload, *_ = result
         return msgspec.json.decode(job_config_payload, type=JobConfig)
     return None
 
 
 @app.get("/<jobid>/config")
 async def show_job_editor(request, jobid: str):
     job = await get_moombox_job_by_id(jobid)
     if not job:
         return "Couldn't find matching job", 404
 
     job_config = _get_saved_jobconf(job.id) or JobConfig.from_moombox_job(job)
 
     return (
         await microdot.jinja.Template("edit.html").render_async(job=job, config=job_config),
         {
             "Content-Type": "text/html; charset=utf-8",
         },
     )
 
 
 @app.post("/<jobid>/config")
 async def apply_job_config(request, jobid: str):
     job = await get_moombox_job_by_id(jobid)
     if not job:
         return "Couldn't find matching job", 404
     jobconf = msgspec.convert(
         {item: request.form[item] for item in request.form}, type=JobConfig
     )
     database = database_ctx.get()
 
     cursor = database.cursor()
     cursor.execute(
         "INSERT INTO job_config VALUES (?, ?) ON CONFLICT(job_id) DO UPDATE SET config=excluded.config",
         (jobid, msgspec.json.encode(jobconf)),
     )
     database.commit()
 
     return (
         await microdot.jinja.Template("edit.html").render_async(job=job, config=jobconf),
         {
             "Content-Type": "text/html; charset=utf-8",
         },
     )
 
 
 @app.get("/<jobid>")
 async def show_job(request, jobid: str):
     return await get_process_job_task(jobid)
 
 
 @app.post("/<jobid>")
 async def upload_job(request, jobid: str):
     try:
         await get_process_job_task(jobid)
         return await microdot.jinja.Template("job/upload_success.html").render_async()
     except torf.TorfError:
         log.exception(f"Failed to upload job {jobid}")
         return await microdot.jinja.Template("job/upload_error.html").render_async()
 
 
 async def job_auto_monitor():
     while True:
         config = config_ctx.get()
 
         async with httpx.AsyncClient() as client:
             for attempt in stamina.retry_context(
                 on=httpx.HTTPError, attempts=None, timeout=None, wait_initial=0.5, wait_max=10.0
             ):
                 with attempt:
                     result = await client.get(f"{config.moombox_url}/status")
 
         jobs = []
         for serialized_job in result.json():
             try:
                 jobs.append(msgspec.convert(serialized_job, type=MoomboxJobInfo))
             except msgspec.ValidationError:
                 pass
 
         database = database_ctx.get()
         cursor = database.cursor()
         configs = {
             job_id: msgspec.json.decode(job_conf, type=JobConfig)
             for job_id, job_conf in cursor.execute("SELECT job_id, config FROM job_config")
         }
         statuses = {
             job_id: bool(completed)
             for job_id, completed in cursor.execute("SELECT job_id, completed FROM job_status")
         }
 
         for job in reversed(jobs):
             if not job.uploadability_state:
                 continue
             job_conf = configs.get(job.id)
             if not job_conf:
                 job_conf = JobConfig.from_moombox_job(job)
             if job.uploadability_state != job_conf.upload:
                 continue
             if not config.autoupload.active:
                 log.info(f"'{job.title}' meets conditions for upload")
                 continue
             if statuses.get(job.id):
                 continue
             log.info(f"'{job.title}' is scheduled for upload")
             get_process_job_task(job.id)
 
         await asyncio.sleep(120)