feat(api): Use FastAPI Error handler

.gitignore

@@ -4,3 +4,4 @@ venv
 data
 *ipynb
 build
+.vscode

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()

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 <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:
-        match mode:
+        case DownloadMode.story:
-            case DownloadMode.story:
+            story_id = download_id
-                story_id = download_id
+        case DownloadMode.part:
-            case DownloadMode.part:
+            story_id = await fetch_story_id(download_id, cookies)
-                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:
+    metadata = await retrieve_story(story_id, cookies)
-            # Invalid ID
+    set_metadata(book, metadata)
-            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>',
-            )
 
-        set_metadata(book, metadata)
+    await set_cover(book, metadata, cookies=cookies)
-        await set_cover(book, metadata, cookies=cookies)
 
-        async for title in add_chapters(
+    async for title in add_chapters(
-            book, metadata, download_images=download_images, cookies=cookies
+        book, metadata, download_images=download_images, cookies=cookies
-        ):
+    ):
-            ...
+        ...
 
-        # Book is compiled
+    # Book is compiled
-        temp_file = tempfile.NamedTemporaryFile(
+    temp_file = tempfile.NamedTemporaryFile(
-            suffix=".epub", delete=True
+        suffix=".epub", delete=True
-        )  # Thanks https://stackoverflow.com/a/75398222
+    )  # Thanks https://stackoverflow.com/a/75398222
 
-        # create epub file
+    # create epub file
-        epub.write_epub(temp_file, book, {})
+    epub.write_epub(temp_file, book, {})
 
-        temp_file.file.seek(0)
+    temp_file.file.seek(0)
-        book_data = temp_file.file.read()
+    book_data = temp_file.file.read()
 
-        return StreamingResponse(
+    return StreamingResponse(
-            BytesIO(book_data),
+        BytesIO(book_data),
-            media_type="application/epub+zip",
+        media_type="application/epub+zip",
-            headers={
+        headers={
-                "Content-Disposition": f'attachment; filename="{slugify(metadata["title"])}_{story_id}_{"images" if download_images else ""}.epub"'  # Thanks https://stackoverflow.com/a/72729058
+            "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>',
-            )
 
 
 app.mount("/", StaticFiles(directory=BUILD_PATH), "static")