Skip to content

Commit

Permalink
New handle_exception plugin hook, refs #1770
Browse files Browse the repository at this point in the history
  • Loading branch information
simonw committed Jul 17, 2022
1 parent 8188f55 commit c09c53f
Show file tree
Hide file tree
Showing 10 changed files with 216 additions and 96 deletions.
99 changes: 25 additions & 74 deletions datasette/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
import secrets
import sys
import threading
import traceback
import urllib.parse
from concurrent import futures
from pathlib import Path
Expand All @@ -27,7 +26,7 @@
from jinja2.environment import Template
from jinja2.exceptions import TemplateNotFound

from .views.base import DatasetteError, ureg
from .views.base import ureg
from .views.database import DatabaseDownload, DatabaseView
from .views.index import IndexView
from .views.special import (
Expand All @@ -49,7 +48,6 @@
PrefixedUrlString,
SPATIALITE_FUNCTIONS,
StartupError,
add_cors_headers,
async_call_with_supported_arguments,
await_me_maybe,
call_with_supported_arguments,
Expand Down Expand Up @@ -87,11 +85,6 @@
from .plugins import pm, DEFAULT_PLUGINS, get_plugins
from .version import __version__

try:
import rich
except ImportError:
rich = None

app_root = Path(__file__).parent.parent

# https://github.com/simonw/datasette/issues/283#issuecomment-781591015
Expand Down Expand Up @@ -1274,6 +1267,16 @@ async def route_path(self, scope, receive, send, path):
return
except NotFound as exception:
return await self.handle_404(request, send, exception)
except Forbidden as exception:
# Try the forbidden() plugin hook
for custom_response in pm.hook.forbidden(
datasette=self.ds, request=request, message=exception.args[0]
):
custom_response = await await_me_maybe(custom_response)
assert (
custom_response
), "Default forbidden() hook should have been called"
return await custom_response.asgi_send(send)
except Exception as exception:
return await self.handle_exception(request, send, exception)

Expand Down Expand Up @@ -1372,72 +1375,20 @@ def raise_404(message=""):
await self.handle_exception(request, send, exception or NotFound("404"))

async def handle_exception(self, request, send, exception):
if self.ds.pdb:
import pdb

pdb.post_mortem(exception.__traceback__)

if rich is not None:
rich.get_console().print_exception(show_locals=True)

title = None
if isinstance(exception, Forbidden):
status = 403
info = {}
message = exception.args[0]
# Try the forbidden() plugin hook
for custom_response in pm.hook.forbidden(
datasette=self.ds, request=request, message=message
):
custom_response = await await_me_maybe(custom_response)
if custom_response is not None:
await custom_response.asgi_send(send)
return
elif isinstance(exception, Base400):
status = exception.status
info = {}
message = exception.args[0]
elif isinstance(exception, DatasetteError):
status = exception.status
info = exception.error_dict
message = exception.message
if exception.message_is_html:
message = Markup(message)
title = exception.title
else:
status = 500
info = {}
message = str(exception)
traceback.print_exc()
templates = [f"{status}.html", "error.html"]
info.update(
{
"ok": False,
"error": message,
"status": status,
"title": title,
}
)
headers = {}
if self.ds.cors:
add_cors_headers(headers)
if request.path.split("?")[0].endswith(".json"):
await asgi_send_json(send, info, status=status, headers=headers)
else:
template = self.ds.jinja_env.select_template(templates)
await asgi_send_html(
send,
await template.render_async(
dict(
info,
urls=self.ds.urls,
app_css_hash=self.ds.app_css_hash(),
menu_links=lambda: [],
)
),
status=status,
headers=headers,
)
responses = []
for hook in pm.hook.handle_exception(
datasette=self.ds,
request=request,
exception=exception,
):
response = await await_me_maybe(hook)
if response is not None:
responses.append(response)

assert responses, "Default exception handler should have returned something"
# Even if there are multiple responses use just the first one
response = responses[0]
await response.asgi_send(send)


_cleaner_task_str_re = re.compile(r"\S*site-packages/")
Expand Down
20 changes: 20 additions & 0 deletions datasette/forbidden.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
from os import stat
from datasette import hookimpl, Response


@hookimpl(trylast=True)
def forbidden(datasette, request, message):
async def inner():
return Response.html(
await datasette.render_template(
"error.html",
{
"title": "Forbidden",
"error": message,
},
request=request,
),
status=403,
)

return inner
74 changes: 74 additions & 0 deletions datasette/handle_exception.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
from datasette import hookimpl, Response
from .utils import await_me_maybe, add_cors_headers
from .utils.asgi import (
Base400,
Forbidden,
)
from .views.base import DatasetteError
from markupsafe import Markup
import pdb
import traceback
from .plugins import pm

try:
import rich
except ImportError:
rich = None


@hookimpl(trylast=True)
def handle_exception(datasette, request, exception):
async def inner():
if datasette.pdb:
pdb.post_mortem(exception.__traceback__)

if rich is not None:
rich.get_console().print_exception(show_locals=True)

title = None
if isinstance(exception, Base400):
status = exception.status
info = {}
message = exception.args[0]
elif isinstance(exception, DatasetteError):
status = exception.status
info = exception.error_dict
message = exception.message
if exception.message_is_html:
message = Markup(message)
title = exception.title
else:
status = 500
info = {}
message = str(exception)
traceback.print_exc()
templates = [f"{status}.html", "error.html"]
info.update(
{
"ok": False,
"error": message,
"status": status,
"title": title,
}
)
headers = {}
if datasette.cors:
add_cors_headers(headers)
if request.path.split("?")[0].endswith(".json"):
return Response.json(info, status=status, headers=headers)
else:
template = datasette.jinja_env.select_template(templates)
return Response.html(
await template.render_async(
dict(
info,
urls=datasette.urls,
app_css_hash=datasette.app_css_hash(),
menu_links=lambda: [],
)
),
status=status,
headers=headers,
)

return inner
5 changes: 5 additions & 0 deletions datasette/hookspecs.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,3 +138,8 @@ def database_actions(datasette, actor, database, request):
@hookspec
def skip_csrf(datasette, scope):
"""Mechanism for skipping CSRF checks for certain requests"""


@hookspec
def handle_exception(datasette, request, exception):
"""Handle an uncaught exception. Can return a Response or None."""
2 changes: 2 additions & 0 deletions datasette/plugins.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
"datasette.default_magic_parameters",
"datasette.blob_renderer",
"datasette.default_menu_links",
"datasette.handle_exception",
"datasette.forbidden",
)

pm = pluggy.PluginManager("datasette")
Expand Down

0 comments on commit c09c53f

Please sign in to comment.