diff --git a/redirect/app.py b/redirect/app.py
index f756c7d..be3b494 100644
--- a/redirect/app.py
+++ b/redirect/app.py
@@ -1,12 +1,12 @@
from typing import Annotated
import django
-from fastapi import FastAPI, Depends, HTTPException, Header
-from fastapi.responses import RedirectResponse
+from fastapi import FastAPI, Depends, Header
+from fastapi.responses import RedirectResponse, HTMLResponse
from sqlalchemy.orm import Session
from starlette.requests import Request
-from redirect.db.curd import get_link_by_slug
+from redirect.db.curd import get_link_by_slug, LinkNotFoundException
from redirect.db.dependency import get_db
@@ -17,6 +17,32 @@
from akarpov.tools.shortener.tasks import save_view_meta # noqa: This has to be imported strictly AFTER django setup
+@app.exception_handler(LinkNotFoundException)
+async def unicorn_exception_handler(request: Request, exc: LinkNotFoundException):
+ return HTMLResponse(
+ status_code=404,
+ # language=HTML
+ content="""
+
+
+ Unknown Link
+
+
+
+ Such link doesn't exist or has been revoked
+
+
+ """,
+ )
+
+
@app.get("/{slug}")
def redirect(slug: str, request: Request, db: Session = Depends(get_db), user_agent: Annotated[str | None, Header()] = None) -> RedirectResponse:
"""Main route that redirects to a page based on the slug."""
@@ -24,8 +50,6 @@ def redirect(slug: str, request: Request, db: Session = Depends(get_db), user_ag
return RedirectResponse(url=f'/tools/shortener/p/{slug.replace("+", "")}')
link = get_link_by_slug(db, slug)
- if link is None:
- raise HTTPException(status_code=404, detail="Unknown Short Link")
save_view_meta.apply_async(
kwargs={
diff --git a/redirect/db/curd.py b/redirect/db/curd.py
index 46af64e..fa78259 100644
--- a/redirect/db/curd.py
+++ b/redirect/db/curd.py
@@ -1,11 +1,14 @@
-from typing import Optional
-
from sqlalchemy.orm import Session
from redirect.db import models
from redirect.util import slug_to_link_id
+class LinkNotFoundException(Exception):
+ def __init__(self, slug: str):
+ self.slug = slug
+
+
def get_link(db: Session, link_id: int):
"""Gets short link metdata by its id.
@@ -19,19 +22,20 @@ def get_link(db: Session, link_id: int):
_get_link_by_slug_cache = {}
-def get_link_by_slug(db: Session, slug: str) -> Optional[tuple[int, str]]:
+def get_link_by_slug(db: Session, slug: str) -> tuple[int, str]:
"""Converts slug to id and gets related link target.
:param db Open connection to a database
:param slug of the link
+ :raise LinkNotFoundException if link is not found
:return target link id & url"""
if slug in _get_link_by_slug_cache:
return _get_link_by_slug_cache[slug]
link = get_link(db, slug_to_link_id(slug))
if link is None or not link.enabled:
- return None
+ raise LinkNotFoundException(slug)
_get_link_by_slug_cache[slug] = (link.id, link.source)
return link.id, link.source