41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
|
from typing import Any, Callable
|
||
|
|
|
||
|
|
from ariadne import SchemaDirectiveVisitor
|
||
|
|
from graphql import GraphQLError
|
||
|
|
|
||
|
|
|
||
|
|
def _wrap_resolver(field: Any, check: Callable[[Any], None]) -> None:
|
||
|
|
original = field.resolve
|
||
|
|
|
||
|
|
def resolve(source: Any, info: Any, **kwargs: Any) -> Any:
|
||
|
|
check(info.context)
|
||
|
|
return original(source, info, **kwargs)
|
||
|
|
|
||
|
|
field.resolve = resolve
|
||
|
|
|
||
|
|
|
||
|
|
# при применении к полю оборачивает оригинальный резолвер:
|
||
|
|
# если в контексте нет `user`, бросает исключение с сообщением «Authentication required»;
|
||
|
|
class IsAuthenticatedDirective(SchemaDirectiveVisitor):
|
||
|
|
def visit_field_definition(self, field: Any, object_type: Any) -> Any:
|
||
|
|
def check(context: Any) -> None:
|
||
|
|
if not context or not context.get("user"):
|
||
|
|
raise GraphQLError("Authentication required")
|
||
|
|
|
||
|
|
_wrap_resolver(field, check)
|
||
|
|
return field
|
||
|
|
|
||
|
|
|
||
|
|
# аналогично проверяет, что `user.is_admin == True`, иначе «Admin only».
|
||
|
|
class IsAdminDirective(SchemaDirectiveVisitor):
|
||
|
|
def visit_field_definition(self, field: Any, object_type: Any) -> Any:
|
||
|
|
def check(context: Any) -> None:
|
||
|
|
user = context.get("user") if context else None
|
||
|
|
if user is None:
|
||
|
|
raise GraphQLError("Authentication required")
|
||
|
|
if not user.is_admin:
|
||
|
|
raise GraphQLError("Admin only")
|
||
|
|
|
||
|
|
_wrap_resolver(field, check)
|
||
|
|
return field
|