From f9e27689e3818e4de354b963a9b69778229ffbc7 Mon Sep 17 00:00:00 2001 From: TheOnlyWayUp Date: Thu, 28 Nov 2024 18:23:52 +0000 Subject: [PATCH] feat(api): Use FastAPI Error handler --- .gitignore | 1 + src/api/src/create_book.py | 9 --- src/api/src/main.py | 118 ++++++++++++++++++------------------- 3 files changed, 58 insertions(+), 70 deletions(-) diff --git a/.gitignore b/.gitignore index 1c01200..20ecf1f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ venv data *ipynb build +.vscode diff --git a/src/api/src/create_book.py b/src/api/src/create_book.py index 1d8d75d..fb6c6c5 100644 --- a/src/api/src/create_book.py +++ b/src/api/src/create_book.py @@ -111,9 +111,6 @@ async def retrieve_story(story_id: int, cookies: Optional[dict] = None) -> dict: async with session.get( f"https://www.wattpad.com/api/v3/stories/{story_id}?fields=tags,id,title,createDate,modifyDate,language(name),description,completed,mature,url,isPaywalled,user(username),parts(id,title),cover" ) as response: - if not response.ok: - if response.status in [404, 400]: - return {} response.raise_for_status() body = await response.json() @@ -132,9 +129,6 @@ async def fetch_part_content(part_id: int, cookies: Optional[dict] = None) -> st async with session.get( f"https://www.wattpad.com/apiv2/?m=storytext&id={part_id}" ) as response: - if not response.ok: - if response.status in [404, 400]: - return "" response.raise_for_status() body = await response.text() @@ -151,9 +145,6 @@ async def fetch_cover(url: str, cookies: Optional[dict] = None) -> bytes: else ClientSession(headers=headers, cookies=cookies) ) as session: # Don't cache requests with Cookies. async with session.get(url) as response: - if not response.ok: - if response.status in [404, 400]: - return bytes() response.raise_for_status() body = await response.read() diff --git a/src/api/src/main.py b/src/api/src/main.py index 21347d5..f1129d6 100644 --- a/src/api/src/main.py +++ b/src/api/src/main.py @@ -1,8 +1,14 @@ +"""WattpadDownloader API Server.""" + from typing import Optional +import tempfile from pathlib import Path +from io import BytesIO from enum import Enum -from fastapi import FastAPI +from aiohttp import ClientResponseError +from fastapi import FastAPI, Request from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse +from fastapi.staticfiles import StaticFiles from ebooklib import epub from create_book import ( retrieve_story, @@ -13,10 +19,6 @@ from create_book import ( wp_get_cookies, fetch_story_id, ) -import tempfile -from io import BytesIO -from fastapi.staticfiles import StaticFiles -from aiohttp import ClientResponseError app = FastAPI() BUILD_PATH = Path(__file__).parent / "build" @@ -37,6 +39,28 @@ def home(): return FileResponse(BUILD_PATH / "index.html") +@app.exception_handler(ClientResponseError) +def download_error_handler(request: Request, exception: ClientResponseError): + match exception.status: + case 400 | 404: + return HTMLResponse( + status_code=404, + content='This story does not exist, or has been deleted. Support is available on the Discord', + ) + case 429: + # Rate-limit by Wattpad + return HTMLResponse( + status_code=429, + content='The website is overloaded. Please try again in a few minutes. Support is available on the Discord', + ) + case _: + # Unhandled error + return HTMLResponse( + status_code=500, + content='Something went wrong. Yell at me on the Discord', + ) + + @app.get("/download/{download_id}") async def handle_download( download_id: int, @@ -45,7 +69,6 @@ async def handle_download( username: Optional[str] = None, password: Optional[str] = None, ): - if username and not password or password and not username: return HTMLResponse( status_code=422, @@ -64,69 +87,42 @@ async def handle_download( else: cookies = None - try: - match mode: - case DownloadMode.story: - story_id = download_id - case DownloadMode.part: - story_id = await fetch_story_id(download_id, cookies) + match mode: + case DownloadMode.story: + story_id = download_id + case DownloadMode.part: + story_id = await fetch_story_id(download_id, cookies) - metadata = await retrieve_story(story_id, cookies) - book = epub.EpubBook() + book = epub.EpubBook() - if not metadata: - # Invalid ID - return HTMLResponse( - status_code=404, - content='The story you tried to download does not exist or has been deleted. Support is available on the Discord', - ) + metadata = await retrieve_story(story_id, cookies) + set_metadata(book, metadata) - set_metadata(book, metadata) - await set_cover(book, metadata, cookies=cookies) + await set_cover(book, metadata, cookies=cookies) - async for title in add_chapters( - book, metadata, download_images=download_images, cookies=cookies - ): - ... + async for title in add_chapters( + book, metadata, download_images=download_images, cookies=cookies + ): + ... - # Book is compiled - temp_file = tempfile.NamedTemporaryFile( - suffix=".epub", delete=True - ) # Thanks https://stackoverflow.com/a/75398222 + # Book is compiled + temp_file = tempfile.NamedTemporaryFile( + suffix=".epub", delete=True + ) # Thanks https://stackoverflow.com/a/75398222 - # create epub file - epub.write_epub(temp_file, book, {}) + # create epub file + epub.write_epub(temp_file, book, {}) - temp_file.file.seek(0) - book_data = temp_file.file.read() + temp_file.file.seek(0) + book_data = temp_file.file.read() - return StreamingResponse( - BytesIO(book_data), - media_type="application/epub+zip", - headers={ - "Content-Disposition": f'attachment; filename="{slugify(metadata["title"])}_{story_id}_{"images" if download_images else ""}.epub"' # Thanks https://stackoverflow.com/a/72729058 - }, - ) - - except ClientResponseError as error: - if error.status == 429: - # Rate-limit by Wattpad - return HTMLResponse( - status_code=429, - content='Unfortunately, the downloader got rate-limited by Wattpad. Please try again later. Support is available on the Discord', - ) - elif error.status == 400: - # Invalid ID - return HTMLResponse( - status_code=404, - content='The story you tried to download does not exist or has been deleted. Support is available on the Discord', - ) - else: - # Unhandled error - return HTMLResponse( - status_code=500, - content='Something went wrong. Support is available on the Discord', - ) + return StreamingResponse( + BytesIO(book_data), + media_type="application/epub+zip", + headers={ + "Content-Disposition": f'attachment; filename="{slugify(metadata["title"])}_{story_id}_{"images" if download_images else ""}.epub"' # Thanks https://stackoverflow.com/a/72729058 + }, + ) app.mount("/", StaticFiles(directory=BUILD_PATH), "static")