mirror of
https://github.com/graphql-python/graphene-django.git
synced 2024-11-22 09:37:07 +03:00
72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
from django.db import connections
|
|
|
|
from promise import Promise
|
|
|
|
from .sql.tracking import unwrap_cursor, wrap_cursor
|
|
from .exception.formating import wrap_exception
|
|
from .types import DjangoDebug
|
|
|
|
|
|
class DjangoDebugContext:
|
|
def __init__(self):
|
|
self.debug_result = None
|
|
self.results = []
|
|
self.object = DjangoDebug(sql=[], exceptions=[])
|
|
self.enable_instrumentation()
|
|
|
|
def get_debug_result(self):
|
|
if not self.debug_result:
|
|
self.debug_result = self.results
|
|
self.results = []
|
|
return self.on_resolve_all_results()
|
|
|
|
def on_resolve_error(self, value):
|
|
if hasattr(self, "object"):
|
|
self.object.exceptions.append(wrap_exception(value))
|
|
return value
|
|
|
|
def on_resolve_all_results(self):
|
|
if self.results:
|
|
self.debug_result = None
|
|
return self.get_debug_result()
|
|
self.disable_instrumentation()
|
|
return self.object
|
|
|
|
def add_result(self, result):
|
|
if self.debug_result:
|
|
self.results.append(result)
|
|
|
|
def enable_instrumentation(self):
|
|
# This is thread-safe because database connections are thread-local.
|
|
for connection in connections.all():
|
|
wrap_cursor(connection, self)
|
|
|
|
def disable_instrumentation(self):
|
|
for connection in connections.all():
|
|
unwrap_cursor(connection)
|
|
|
|
|
|
class DjangoDebugMiddleware:
|
|
def resolve(self, next, root, info, **args):
|
|
context = info.context
|
|
django_debug = getattr(context, "django_debug", None)
|
|
if not django_debug:
|
|
if context is None:
|
|
raise Exception("DjangoDebug cannot be executed in None contexts")
|
|
try:
|
|
context.django_debug = DjangoDebugContext()
|
|
except Exception:
|
|
raise Exception(
|
|
"DjangoDebug need the context to be writable, context received: {}.".format(
|
|
context.__class__.__name__
|
|
)
|
|
)
|
|
if info.schema.get_type("DjangoDebug") == info.return_type:
|
|
return context.django_debug.get_debug_result()
|
|
try:
|
|
result = next(root, info, **args)
|
|
except Exception as e:
|
|
return context.django_debug.on_resolve_error(e)
|
|
context.django_debug.add_result(result)
|
|
return result
|