feat(api): Use FastAPI Error handler

2024-11-28 18:23:52 +00:00
parent 308afde25f
commit f9e27689e3
3 changed files with 58 additions and 70 deletions
@@ -4,3 +4,4 @@ venv
 data
 *ipynb
 build
+.vscode
@@ -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()
@@ -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 <a href="https://discord.gg/P9RHC4KCwd" target="_blank">Discord</a>',
+            )
+        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 <a href="https://discord.gg/P9RHC4KCwd" target="_blank">Discord</a>',
+            )
+        case _:
+            # Unhandled error
+            return HTMLResponse(
+                status_code=500,
+                content='Something went wrong. Yell at me on the <a href="https://discord.gg/P9RHC4KCwd" target="_blank">Discord</a>',
+            )
+
+
@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 <a href="https://discord.gg/P9RHC4KCwd" target="_blank">Discord</a>',
-            )
+    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 <a href="https://discord.gg/P9RHC4KCwd" target="_blank">Discord</a>',
-            )
-        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 <a href="https://discord.gg/P9RHC4KCwd" target="_blank">Discord</a>',
-            )
-        else:
-            # Unhandled error
-            return HTMLResponse(
-                status_code=500,
-                content='Something went wrong. Support is available on the <a href="https://discord.gg/P9RHC4KCwd" target="_blank">Discord</a>',
-            )
+    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")