From b05fe47914df8dbc89711c87d717ab427a7d6cc9 Mon Sep 17 00:00:00 2001
From: TheOnlyWayUp <hi@towu.dev>
Date: Fri, 6 Dec 2024 07:56:08 +0000
Subject: [PATCH] feat(api): Errors are raised faster, add Exception classes

---
 src/api/src/create_book.py | 50 +++++++++++++++++++++++++++++++++-----
 src/api/src/main.py        |  9 +++++++
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/src/api/src/create_book.py b/src/api/src/create_book.py
index 93356a6..074474f 100644
--- a/src/api/src/create_book.py
+++ b/src/api/src/create_book.py
@@ -1,6 +1,7 @@
 from typing import List, Optional, Tuple
 from typing_extensions import TypedDict
 import re
+import json
 import unicodedata
 import logging
 from os import environ
@@ -195,6 +196,22 @@ class Story(TypedDict):
 
 story_ta = TypeAdapter(Story)
 
+# --- Exceptions --- #
+
+
+class WattpadError(Exception):
+    """Base Exception class for Wattpad related errors."""
+
+
+class StoryNotFoundError(WattpadError):
+    """Display the "This story was not found" error to the user."""
+
+    ...
+
+
+class PartNotFoundError(StoryNotFoundError): ...
+
+
 # --- API Calls --- #
 
 
@@ -210,10 +227,16 @@ async def fetch_story_from_partId(
             async with session.get(
                 f"https://www.wattpad.com/api/v3/story_parts/{part_id}?fields=groupId,group(tags,id,title,createDate,modifyDate,language(name),description,completed,mature,url,isPaywalled,user(username),parts(id,title),cover)"
             ) as response:
-                response.raise_for_status()
-
                 body = await response.json()
 
+                if response.status == 400:
+                    match body.get("error_code"):
+                        case 1020:  # "Story part not found"
+                            logger.info(f"{part_id=} not found on Wattpad, returning.")
+                            raise PartNotFoundError()
+
+                response.raise_for_status()
+
         return str(body["groupId"]), story_ta.validate_python(body["group"])
 
 
@@ -227,10 +250,16 @@ async def retrieve_story(story_id: int, cookies: Optional[dict] = None) -> Story
             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:
-                response.raise_for_status()
-
                 body = await response.json()
 
+                if response.status == 400:
+                    match body.get("error_code"):
+                        case 1017:  # "Story not found"
+                            logger.info(f"{story_id=} not found on Wattpad, returning.")
+                            raise StoryNotFoundError()
+
+                response.raise_for_status()
+
         return story_ta.validate_python(body)
 
 
@@ -244,10 +273,19 @@ 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:
-                response.raise_for_status()
-
                 body = await response.text()
 
+                if response.status == 400:
+                    data = json.loads(body)
+                    match data.get("code"):
+                        case 463:  # ""Could not find any parts for that story""
+                            logger.info(
+                                f"{part_id=} for text not found on Wattpad, returning."
+                            )
+                            raise PartNotFoundError()
+
+                response.raise_for_status()
+
         return body
 
 
diff --git a/src/api/src/main.py b/src/api/src/main.py
index c4d1824..b6ff247 100644
--- a/src/api/src/main.py
+++ b/src/api/src/main.py
@@ -104,6 +104,15 @@ def download_error_handler(request: Request, exception: ClientResponseError):
             )
 
 
+@app.exception_handler(WattpadError)
+def download_wp_error_handler(request: Request, exception: WattpadError):
+    if isinstance(exception, StoryNotFoundError):
+        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>',
+        )
+
+
 @app.get("/download/{download_id}")
 async def handle_download(
     download_id: int,